use mdarray::{view, tensor, Slice, Dim, Const, expr::Expression};
use mdarray as md;

// Indexing convention: C_ij <- A_ik * B_kj
fn matmul<D0: Dim, D1: Dim, D2: Dim>(
    a: &Slice<f64, (D0, D1)>,
    b: &Slice<f64, (D1, D2)>,
    c: &mut Slice<f64, (D0, D2)>) {
    for (mut ci, ai) in c.rows_mut().zip(a.rows()) {
        for (aik, bk) in ai.zip(b.rows()) {
            for (cij, bkj) in ci.expr_mut().zip(bk) {
                *cij = aik.mul_add(*bkj, *cij);
            }
        }
    }
}

fn main() {
    let a = view![[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]];
    let b = md::Array::<f64, (Const::<2>, Const::<2>)>::from_fn(
        |i| if i[0] == 0 && i[1] == 0 { 0.0 } else { 1.0 }
    );
    let b = b.reshape((Const::<2>, !0));

    // .into_dyn() replaces .reshape(DynRank::from_dims(&[2, 3])).into():
    let mut c: mdarray::Tensor<f64> = tensor![[0.0; 2]; 3].into_dyn();

    dbg!(std::any::type_name_of_val(&a));
    dbg!(std::any::type_name_of_val(&b));
    dbg!(std::any::type_name_of_val(&c));

    matmul(&a.reshape((3, Const::<2>)), &b, &mut c.remap_mut());

    assert_eq!(c, view![[4.0, 5.0], [5.0, 7.0], [6.0, 9.0]]);

    // slice
    let d = a.view(1, ..);
    dbg!(std::any::type_name_of_val(&d));

    // index
    let e = c[4];
    dbg!(std::any::type_name_of_val(&e));

    // permute
    let f = c.permute([1, 0]);
    dbg!(std::any::type_name_of_val(&f));

    // Arrays with rank > 6:
    let x: Vec<f64> = (0..128).map(|x| x as f64).collect();
    let x: md::DTensor<f64, 1> = x.into();
    let x = x.into_dyn();
    let x = x.reshape(&[2; 7]);
    let x = x.permute(&[6, 5, 4, 3, 2, 1, 0]);
    dbg!(std::any::type_name_of_val(&x));

    // Indexing
    for i in 0..2 {
        let index = [0, i, 0, 0, i, 0, 0];
        println!("{} {} {} {}", a[i], a[[i, 0]], c[&index[..2]], x[index]);
    }
}