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

struct SumAccu {
    value: C64,
}

impl SumAccu {
    fn new() -> Self {
        Self { value: C64::ZERO }
    }
    fn accumulate(&mut self, x: C64) {
        self.value = self.value + x;
    }
    fn get(&self) -> C64 {
        self.value
    }
}

fn pure_ndarray(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    let intermediate: Array3<C64> = Zip::from(t_a).and(t_b).map_collect(|&a_val, &b_val| {
        if a_val > C64::ZERO {
            a_val * b_val
        } else {
            (C64::ONE + b_val.exp()).ln()
        }
    });
    intermediate.sum_axis(Axis(1)).t().to_owned()
}

fn ndarray_non_custom(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    let intermediate: Array3<C64> = Zip::from(t_a).and(t_b).map_collect(|&a_val, &b_val| {
        if a_val > C64::ZERO {
            a_val * b_val
        } else {
            (C64::ONE + b_val.exp()).ln()
        }
    });
    intermediate
        .map_axis(Axis(1), |row| row.iter().map(|x| *x).sum())
        .t()
        .to_owned()
}

fn tensorism(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    new_ndarray! {for k i => (for j =>
        let elt_a = t_a[i, j, k];
        let elt_b = t_b[i, j, k];
        if elt_a > C64::ZERO {
            elt_a * elt_b
        } else {
            (C64::ONE + elt_b.exp()).ln()
        }).sum::<C64>()
    }
}

fn flatten_factorized(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
    }
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
        panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
    }
    let strides_a = t_a.strides();
    if strides_a != [20isize, 4isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let strides_b = t_b.strides();
    if strides_b != [20isize, 4isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let stride_a0 = strides_a[0];
    let stride_a1 = strides_a[1];
    let stride_a2 = strides_a[2];
    let stride_b0 = strides_b[0];
    let stride_b1 = strides_b[1];
    let stride_b2 = strides_b[2];
    let p_a = t_a.as_ptr();
    let p_b = t_b.as_ptr();
    ::ndarray::Array::<_, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::from_shape_fn(
        (dim_number_0, dim_number_1),
        |(k, i)| {
            let pre_a = (i as isize) * stride_a0 + (k as isize) * stride_a2;
            let pre_b = (i as isize) * stride_b0 + (k as isize) * stride_b2;
            ((0usize..dim_number_2).map(|j| {
                let elt_a: C64 = unsafe { *p_a.offset(pre_a + (j as isize) * stride_a1) };
                let elt_b: C64 = unsafe { *p_b.offset(pre_b + (j as isize) * stride_b1) };
                if elt_a > C64::ZERO {
                    elt_a * elt_b
                } else {
                    (C64::ONE + elt_b.exp()).ln()
                }
            }))
            .sum::<C64>()
        },
    )
}

fn flatten_normal(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
    }
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
        panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
    }
    let strides_a = t_a.strides();
    if strides_a != [20isize, 4isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let strides_b = t_b.strides();
    if strides_b != [20isize, 4isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let stride_a0 = strides_a[0];
    let stride_a1 = strides_a[1];
    let stride_a2 = strides_a[2];
    let stride_b0 = strides_b[0];
    let stride_b1 = strides_b[1];
    let stride_b2 = strides_b[2];
    let p_a = t_a.as_ptr();
    let p_b = t_b.as_ptr();

    let mut result = ::ndarray::Array::<C64, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
        dim_number_0,
        dim_number_1,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut C64;

    for k in 0..dim_number_0 {
        for i in 0..dim_number_1 {
            let value = ((0usize..dim_number_2).map(|j| {
                let elt_a: C64 = unsafe {
                    *p_a.offset(
                        (i as isize) * stride_a0
                            + (k as isize) * stride_a2
                            + (j as isize) * stride_a1,
                    )
                };
                let elt_b: C64 = unsafe {
                    *p_b.offset(
                        (i as isize) * stride_b0
                            + (k as isize) * stride_b2
                            + (j as isize) * stride_b1,
                    )
                };
                if elt_a > C64::ZERO {
                    elt_a * elt_b
                } else {
                    (C64::ONE + elt_b.exp()).ln()
                }
            }))
            .sum::<C64>();
            unsafe {
                ptr_on_result.write(value);
                ptr_on_result = ptr_on_result.add(1);
            }
        }
    }
    unsafe { result.assume_init() }
}

fn flatten_accumulator(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
    }
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
        panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
    }
    let strides_a = t_a.strides();
    if strides_a != [20isize, 4isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let strides_b = t_b.strides();
    if strides_b != [20isize, 4isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let stride_a0 = strides_a[0];
    let stride_a1 = strides_a[1];
    let stride_a2 = strides_a[2];
    let stride_b0 = strides_b[0];
    let stride_b1 = strides_b[1];
    let stride_b2 = strides_b[2];
    let p_a = t_a.as_ptr();
    let p_b = t_b.as_ptr();

    let mut result = ::ndarray::Array::<C64, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
        dim_number_0,
        dim_number_1,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut C64;

    for k in 0..dim_number_0 {
        for i in 0..dim_number_1 {
            let mut accu = SumAccu::new();
            for j in 0..dim_number_2 {
                let elt_a: C64 = unsafe {
                    *p_a.offset(
                        (i as isize) * stride_a0
                            + (k as isize) * stride_a2
                            + (j as isize) * stride_a1,
                    )
                };
                let elt_b: C64 = unsafe {
                    *p_b.offset(
                        (i as isize) * stride_b0
                            + (k as isize) * stride_b2
                            + (j as isize) * stride_b1,
                    )
                };
                accu.accumulate(if elt_a > C64::ZERO {
                    elt_a * elt_b
                } else {
                    (C64::ONE + elt_b.exp()).ln()
                });
            }
            unsafe {
                ptr_on_result.write(accu.get());
                ptr_on_result = ptr_on_result.add(1);
            }
        }
    }
    unsafe { result.assume_init() }
}

fn flatten_accumulator_factorized(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
    }
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
        panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
    }
    let strides_a = t_a.strides();
    if strides_a != [20isize, 4isize, 1isize] {
        panic!("t_a has unexpected memory layout {:?}", strides_a);
    }
    let strides_b = t_b.strides();
    if strides_b != [20isize, 4isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let stride_a0 = strides_a[0];
    let stride_a1 = strides_a[1];
    let stride_a2 = strides_a[2];
    let stride_b0 = strides_b[0];
    let stride_b1 = strides_b[1];
    let stride_b2 = strides_b[2];
    let p_a = t_a.as_ptr();
    let p_b = t_b.as_ptr();

    let mut result = ::ndarray::Array::<C64, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
        dim_number_0,
        dim_number_1,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut C64;

    for k in 0..dim_number_0 {
        for i in 0..dim_number_1 {
            let mut accu = SumAccu::new();
            let pre_a = (i as isize) * stride_a0 + (k as isize) * stride_a2;
            let pre_b = (i as isize) * stride_b0 + (k as isize) * stride_b2;
            for j in 0..dim_number_2 {
                let elt_a: C64 = unsafe { *p_a.offset(pre_a + (j as isize) * stride_a1) };
                let elt_b: C64 = unsafe { *p_b.offset(pre_b + (j as isize) * stride_b1) };
                accu.accumulate(if elt_a > C64::ZERO {
                    elt_a * elt_b
                } else {
                    (C64::ONE + elt_b.exp()).ln()
                });
            }
            unsafe {
                ptr_on_result.write(accu.get());
                ptr_on_result = ptr_on_result.add(1);
            }
        }
    }
    unsafe { result.assume_init() }
}

// result[k, i] = SUM_{j} (if A[i,j,k]>0 then A[i,j,k]*B[i,j,k] else ln(1+exp(B[i,j,k])))
fn benchmark(c: &mut Criterion) {
    let shape = (8, 5, 4);
    let tensor1 = Array3::<C64>::from_shape_fn(shape, |(i, j, k)| {
        let i = C64::from_usize(i) * PI;
        let j = C64::from_usize(j) * PI + C64::I;
        let k = C64::from_usize(k) * PI;
        (i * k + 3.0).sin() * 0.9 - (j / 2.0 + 1.44).sin() * 1.5 + (k / 3.0 + 4.9 - j).cos() * 2.7
    });
    let tensor2 = Array3::<C64>::from_shape_fn(shape, |(i, j, k)| {
        let i = C64::from_usize(i) * PI;
        let j = C64::from_usize(j) * PI;
        let k = C64::from_usize(k) * PI + C64::I;
        (j * i + 2.9).sin() * 1.1 - (k / 1.9 + 1.54).sin() * 1.6 + (j / 3.1 + 4.8 - k).cos() * 2.6
    });
    let result_ndarray = pure_ndarray(&tensor1, &tensor2);
    let result_ndarray_non_custom = ndarray_non_custom(&tensor1, &tensor2);
    let result_tensorism = tensorism(&tensor1, &tensor2);
    let result_flatten = flatten_factorized(&tensor1, &tensor2);
    let result_flatten_var = flatten_normal(&tensor1, &tensor2);
    assert_eq!(result_ndarray, result_tensorism);
    assert_eq!(result_ndarray_non_custom, result_tensorism);
    assert_eq!(result_flatten, result_tensorism);
    assert_eq!(result_flatten_var, result_tensorism);

    c.bench_function("Conditional value - Ndarray", |b| {
        b.iter(|| pure_ndarray(black_box(&tensor1), black_box(&tensor2)))
    });
    c.bench_function("Conditional value - Ndarray Non Custom", |b| {
        b.iter(|| ndarray_non_custom(black_box(&tensor1), black_box(&tensor2)))
    });
    c.bench_function("Conditional value - Flatten normal", |b| {
        b.iter(|| flatten_normal(black_box(&tensor1), black_box(&tensor2)))
    });
    c.bench_function("Conditional value - Flatten factorized", |b| {
        b.iter(|| flatten_factorized(black_box(&tensor1), black_box(&tensor2)))
    });
    c.bench_function("Conditional value - Tensorism", |b| {
        b.iter(|| tensorism(black_box(&tensor1), black_box(&tensor2)))
    });
    c.bench_function("Conditional value - Flatten accumulator", |b| {
        b.iter(|| flatten_accumulator(black_box(&tensor1), black_box(&tensor2)))
    });
    c.bench_function("Conditional value - Flatten accumulator factorized", |b| {
        b.iter(|| flatten_accumulator_factorized(black_box(&tensor1), black_box(&tensor2)))
    });
}

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