use criterion::{Criterion, criterion_group, criterion_main};
use ndarray::{Array1, Array2, Array3, Axis, Zip};
use std::hint::black_box;
use tensorism::new_ndarray;
use utils::C64;

fn tensorism(
    t_a: &Array2<f64>,
    t_b: &Array1<C64>,
    t_c: &Array2<f64>,
    t_d: &Array1<i32>,
) -> Array3<f64> {
    new_ndarray! {
        for i j k => t_a[j, i] * t_c[i, k] + t_b[k].re() * t_b[k].re() + (t_d[j] as f64).sin()
    }
}

fn flatten_low(
    t_a: &Array2<f64>,
    t_b: &Array1<C64>,
    t_c: &Array2<f64>,
    t_d: &Array1<i32>,
) -> Array3<f64> {
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_c).0 {
        panic!("Dimensions are not matching between t_a[_, i] and t_c[i, _]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_d) {
        panic!("Dimensions are not matching between t_a[j, _] and t_d[j]");
    }
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_c).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b) {
        panic!("Dimensions are not matching between t_c[_, k] and t_b[k]");
    }
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b) {
        panic!("Dimensions are not matching between t_c[_, k] and t_b[k]");
    }
    let strides_a = t_a.strides();
    if strides_a != [dim_number_0 as isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let p_a = t_a.as_ptr();
    let strides_b = t_b.strides();
    if strides_b != [1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let p_b = t_b.as_ptr();
    let strides_c = t_c.strides();
    if strides_c != [dim_number_2 as isize, 1isize] {
        panic!("t_c has unexpected memory layout {:?}", strides_c);
    }
    let p_c = t_c.as_ptr();
    let strides_d = t_d.strides();
    if strides_d != [1isize] {
        panic!("t_d has unexpected memory layout {:?}", strides_d);
    }
    let p_d = t_d.as_ptr();

    let mut result = ::ndarray::Array::<f64, ::ndarray::Dim<[::ndarray::Ix; 3usize]>>::uninit((
        dim_number_0,
        dim_number_1,
        dim_number_2,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut f64;
    for i in 0..dim_number_0 {
        for j in 0..dim_number_1 {
            for k in 0..dim_number_2 {
                let value = (unsafe { *p_a.add(j * dim_number_0 + i) })
                    * (unsafe { *p_c.add(i * dim_number_2 + k) })
                    + (unsafe { *p_b.add(k) }).re() * (unsafe { *p_b.add(k) }).re()
                    + ((unsafe { *p_d.add(j) }) as f64).sin();
                unsafe {
                    ptr_on_result.write(value);
                    ptr_on_result = ptr_on_result.add(1);
                }
            }
        }
    }
    unsafe { result.assume_init() }
}

fn flatten_high(
    t_a: &Array2<f64>,
    t_b: &Array1<C64>,
    t_c: &Array2<f64>,
    t_d: &Array1<i32>,
) -> Array3<f64> {
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_c).0 {
        panic!("Dimensions are not matching between t_a[_, i] and t_c[i, _]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_d) {
        panic!("Dimensions are not matching between t_a[j, _] and t_d[j]");
    }
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_c).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b) {
        panic!("Dimensions are not matching between t_c[_, k] and t_b[k]");
    }
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b) {
        panic!("Dimensions are not matching between t_c[_, k] and t_b[k]");
    }
    let strides_a = t_a.strides();
    if strides_a != [dim_number_0 as isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let p_a = t_a.as_ptr();
    let strides_b = t_b.strides();
    if strides_b != [1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let p_b = t_b.as_ptr();
    let strides_c = t_c.strides();
    if strides_c != [dim_number_2 as isize, 1isize] {
        panic!("t_c has unexpected memory layout {:?}", strides_c);
    }
    let p_c = t_c.as_ptr();
    let strides_d = t_d.strides();
    if strides_d != [1isize] {
        panic!("t_d has unexpected memory layout {:?}", strides_d);
    }
    let p_d = t_d.as_ptr();

    let mut result = ::ndarray::Array::<f64, ::ndarray::Dim<[::ndarray::Ix; 3usize]>>::uninit((
        dim_number_0,
        dim_number_1,
        dim_number_2,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut f64;
    for i in 0..dim_number_0 {
        let i_c = i * dim_number_2;
        for j in 0..dim_number_1 {
            let i_a = j * dim_number_0 + i;
            for k in 0..dim_number_2 {
                let value = (unsafe { *p_a.add(i_a) }) * (unsafe { *p_c.add(i_c + k) })
                    + (unsafe { *p_b.add(k) }).re() * (unsafe { *p_b.add(k) }).re()
                    + ((unsafe { *p_d.add(j) }) as f64).sin();
                unsafe {
                    ptr_on_result.write(value);
                    ptr_on_result = ptr_on_result.add(1);
                }
            }
        }
    }
    unsafe { result.assume_init() }
}

fn flatten_var(
    t_a: &Array2<f64>,
    t_b: &Array1<C64>,
    t_c: &Array2<f64>,
    t_d: &Array1<i32>,
) -> Array3<f64> {
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_c).0 {
        panic!("Dimensions are not matching between t_a[_, i] and t_c[i, _]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_d) {
        panic!("Dimensions are not matching between t_a[j, _] and t_d[j]");
    }
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_c).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b) {
        panic!("Dimensions are not matching between t_c[_, k] and t_b[k]");
    }
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b) {
        panic!("Dimensions are not matching between t_c[_, k] and t_b[k]");
    }
    let strides_a = t_a.strides();
    if strides_a != [dim_number_0 as isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let p_a = t_a.as_ptr();
    let strides_b = t_b.strides();
    if strides_b != [1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let p_b = t_b.as_ptr();
    let strides_c = t_c.strides();
    if strides_c != [dim_number_2 as isize, 1isize] {
        panic!("t_c has unexpected memory layout {:?}", strides_c);
    }
    let p_c = t_c.as_ptr();
    let strides_d = t_d.strides();
    if strides_d != [1isize] {
        panic!("t_d has unexpected memory layout {:?}", strides_d);
    }
    let p_d = t_d.as_ptr();

    let mut result = ::ndarray::Array::<f64, ::ndarray::Dim<[::ndarray::Ix; 3usize]>>::uninit((
        dim_number_0,
        dim_number_1,
        dim_number_2,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut f64;
    let mut i_1_a: usize = 0;
    let mut i_1_c: usize = 0;
    for _i in 0..dim_number_0 {
        let mut i_2_a: usize = i_1_a;
        let mut i_2_d: usize = 0;
        for _j in 0..dim_number_1 {
            let mut i_3_b: usize = 0;
            let mut i_3_c: usize = i_1_c;
            for _k in 0..dim_number_2 {
                let value = (unsafe { *p_a.add(i_2_a) }) * (unsafe { *p_c.add(i_3_c) })
                    + (unsafe { *p_b.add(i_3_b) }).re() * (unsafe { *p_b.add(i_3_b) }).re()
                    + ((unsafe { *p_d.add(i_2_d) }) as f64).sin();
                unsafe {
                    ptr_on_result.write(value);
                    ptr_on_result = ptr_on_result.add(1);
                }
                i_3_b += 1;
                i_3_c += 1;
            }
            i_2_a += dim_number_0;
            i_2_d += 1;
        }
        i_1_a += 1;
        i_1_c += dim_number_2;
    }
    unsafe { result.assume_init() }
}

fn pure_ndarray(
    t_a: &Array2<f64>,
    t_b: &Array1<C64>,
    t_c: &Array2<f64>,
    t_d: &Array1<i32>,
) -> Array3<f64> {
    let (ni, nj, nk) = (128, 64, 32);

    let mut r = Array3::<_>::uninit((ni, nj, nk));

    let t_a_perm = t_a.view().permuted_axes([1, 0]);
    let t_a_ins = t_a_perm.insert_axis(Axis(2));
    let t_a_v = t_a_ins.broadcast((ni, nj, nk)).unwrap();

    let t_b_ins0 = t_b.view().insert_axis(Axis(0));
    let t_b_ins1 = t_b_ins0.insert_axis(Axis(1));
    let t_b_v = t_b_ins1.broadcast((ni, nj, nk)).unwrap();

    let t_c_ins = t_c.view().insert_axis(Axis(1));
    let t_c_v = t_c_ins.broadcast((ni, nj, nk)).unwrap();

    let t_d_ins0 = t_d.view().insert_axis(Axis(0));
    let t_d_ins1 = t_d_ins0.insert_axis(Axis(2));
    let t_d_v = t_d_ins1.broadcast((ni, nj, nk)).unwrap();

    Zip::from(r.view_mut())
        .and(t_a_v)
        .and(t_b_v)
        .and(t_c_v)
        .and(t_d_v)
        .for_each(|r_cell, &a, &b, &c, &d| {
            r_cell.write(a * c + b.re() * b.re() + (d as f64).sin());
        });

    unsafe { r.assume_init() }
}

fn benchmark(c: &mut Criterion) {
    let (ni, nj, nk) = (128, 64, 32);
    let t_a: Array2<f64> = Array2::from_shape_fn((nj, ni), |(j, i)| {
        1.7 * ((i as f64) - (j as f64) + 0.5).cos()
    });
    let t_b: Array1<C64> =
        Array1::from_shape_fn(nk, |k| C64::from_f64(3.5 * (k as f64 + 2.0).sin()));
    let t_c: Array2<f64> = Array2::from_shape_fn((ni, nk), |(i, k)| ((i * k) as f64 + 10.0).ln());
    let t_d: Array1<i32> = Array1::from_shape_fn(nj, |j| j as i32);

    let result_ndarray = pure_ndarray(&t_a, &t_b, &t_c, &t_d);
    let result_tensorism = tensorism(&t_a, &t_b, &t_c, &t_d);
    let result_flatten_low = flatten_low(&t_a, &t_b, &t_c, &t_d);
    let result_flatten_var = flatten_var(&t_a, &t_b, &t_c, &t_d);
    let result_flatten_high = flatten_high(&t_a, &t_b, &t_c, &t_d);
    assert_eq!(result_tensorism, result_flatten_low);
    assert_eq!(result_tensorism, result_flatten_var);
    assert_eq!(result_tensorism, result_flatten_high);
    assert_eq!(result_ndarray, result_tensorism);

    c.bench_function("Mixed array - Ndarray", |b| {
        b.iter(|| {
            pure_ndarray(
                black_box(&t_a),
                black_box(&t_b),
                black_box(&t_c),
                black_box(&t_d),
            )
        })
    });
    c.bench_function("Mixed array - Tensorism", |b| {
        b.iter(|| {
            tensorism(
                black_box(&t_a),
                black_box(&t_b),
                black_box(&t_c),
                black_box(&t_d),
            )
        })
    });
    c.bench_function("Mixed array - Flatten Low", |b| {
        b.iter(|| {
            flatten_low(
                black_box(&t_a),
                black_box(&t_b),
                black_box(&t_c),
                black_box(&t_d),
            )
        })
    });
    c.bench_function("Mixed array - Flatten Var", |b| {
        b.iter(|| {
            flatten_var(
                black_box(&t_a),
                black_box(&t_b),
                black_box(&t_c),
                black_box(&t_d),
            )
        })
    });
    c.bench_function("Mixed array - Flatten High", |b| {
        b.iter(|| {
            flatten_high(
                black_box(&t_a),
                black_box(&t_b),
                black_box(&t_c),
                black_box(&t_d),
            )
        })
    });
}

criterion_group!(benches, benchmark);
criterion_main!(benches);
