use criterion::{Criterion, black_box, criterion_group, criterion_main};
use ndarray::{Array2, Array3, Zip, s};
use std::collections::HashSet;
use tensorism::new_ndarray;

struct AnyAccu {
    value: bool,
}

impl AnyAccu {
    fn new() -> Self {
        Self { value: false }
    }
    #[inline(always)]
    fn get(&self) -> bool {
        self.value
    }
}

fn tensorism(t_a: &Array3<HashSet<usize>>, t_b: &Array2<usize>) -> Array2<bool> {
    new_ndarray! {
        for i j => (for k => t_a[i, k, j].contains(&t_b[j, i])).any(|b| b)
    }
}

fn flatten(t_a: &Array3<HashSet<usize>>, t_b: &Array2<usize>) -> Array2<bool> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[_, i]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[_, _, j] and t_b[j, _]");
    }

    let strides_a = t_a.strides();
    if strides_a != [91isize, 7isize, 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 != [11isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let p_b = t_b.as_ptr();
    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 mut result = ::ndarray::Array::<bool, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
        dim_number_0,
        dim_number_1,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut bool;
    for i in 0..dim_number_0 {
        for j in 0..dim_number_1 {
            let value: bool = ((0usize..dim_number_2).map(|k| {
                (unsafe {
                    &*p_a.offset(
                        stride_a0 * (i as isize)
                            + stride_a1 * (k as isize)
                            + stride_a2 * (j as isize),
                    )
                })
                .contains(
                    &(unsafe { *p_b.offset(stride_b0 * (j as isize) + stride_b1 * (i as isize)) }),
                )
            }))
            .any(|b| b);
            unsafe {
                ptr_on_result.write(value);
                ptr_on_result = ptr_on_result.add(1);
            }
        }
    }
    unsafe { result.assume_init() }
}

fn flatten_accumulator(t_a: &Array3<HashSet<usize>>, t_b: &Array2<usize>) -> Array2<bool> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[_, i]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[_, _, j] and t_b[j, _]");
    }

    let strides_a = t_a.strides();
    if strides_a != [91isize, 7isize, 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 != [11isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let p_b = t_b.as_ptr();
    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 mut result = ::ndarray::Array::<bool, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
        dim_number_0,
        dim_number_1,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut bool;
    for i in 0..dim_number_0 {
        for j in 0..dim_number_1 {
            let mut accu = AnyAccu::new();
            for k in 0..dim_number_2 {
                accu.value |= (unsafe {
                    &*p_a.offset(
                        stride_a0 * (i as isize)
                            + stride_a1 * (k as isize)
                            + stride_a2 * (j as isize),
                    )
                })
                .contains(
                    &(unsafe { *p_b.offset(stride_b0 * (j as isize) + stride_b1 * (i as isize)) }),
                );
            }
            unsafe {
                ptr_on_result.write(accu.get());
                ptr_on_result = ptr_on_result.add(1);
            }
        }
    }
    unsafe { result.assume_init() }
}

fn flatten_factorized(t_a: &Array3<HashSet<usize>>, t_b: &Array2<usize>) -> Array2<bool> {
    let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
    let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
    if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
        panic!("Dimensions are not matching between t_a[i, _, _] and t_b[_, i]");
    }
    let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
    if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
        panic!("Dimensions are not matching between t_a[_, _, j] and t_b[j, _]");
    }

    let strides_a = t_a.strides();
    if strides_a != [91isize, 7isize, 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 != [11isize, 1isize] {
        panic!("t_b has unexpected memory layout {:?}", strides_b);
    }
    let p_b = t_b.as_ptr();
    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 mut result = ::ndarray::Array::<bool, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
        dim_number_0,
        dim_number_1,
    ));
    let mut ptr_on_result = result.as_mut_ptr() as *mut bool;
    for i in 0..dim_number_0 {
        for j in 0..dim_number_1 {
            let i_a = stride_a0 * (i as isize) + stride_a2 * (j as isize);
            let i_b = stride_b0 * (j as isize) + stride_b1 * (i as isize);
            let value: bool = ((0usize..dim_number_2).map(|k| {
                (unsafe { &*p_a.offset(i_a + stride_a1 * (k as isize)) })
                    .contains(&(unsafe { *p_b.offset(i_b) }))
            }))
            .any(|b| b);
            unsafe {
                ptr_on_result.write(value);
                ptr_on_result = ptr_on_result.add(1);
            }
        }
    }
    unsafe { result.assume_init() }
}

fn pure_ndarray(t_a: &Array3<HashSet<usize>>, t_b: &Array2<usize>) -> Array2<bool> {
    let i_max = t_a.dim().0;
    let j_max = t_a.dim().2;

    // Création du tenseur de résultat C[i, j] initialisé à false
    let mut result = Array2::<bool>::from_elem((i_max, j_max), false);

    // Zip permet d'itérer sur le résultat C[i, j] et les vues t_a[i, :, j] et t_b[j, i]
    Zip::indexed(&mut result).for_each(|(i, j), output_val| {
        // 1. Récupération de la valeur scalaire de référence de t_b[j, i]
        let b_val_ref = &t_b[(j, i)]; // Note: t_b est indexé (j, i)

        // 2. Création d'une vue sur l'axe de réduction k (l'axe 1 de t_a)
        // t_a_slice est la colonne t_a[i, :, j]
        let t_a_slice = t_a.slice(s![i, .., j]);

        // 3. Boucle d'agrégation manuelle (ANY)
        // Itérer sur tous les éléments de l'axe k
        for set in t_a_slice.iter() {
            if set.contains(b_val_ref) {
                *output_val = true;
                break;
            }
        }
    });

    result
}

// result[i, j] = ANY_{k} t_a[i, k, j].contains(t_b[j, i])
fn benchmark(c: &mut Criterion) {
    let (di, dj, dk) = (11, 7, 13);

    let shape_a = (di, dk, dj);
    let t_a = Array3::<HashSet<usize>>::from_shape_fn(shape_a, |(i, k, j)| {
        let v = (i * 5 + j * 3 + 11) % 17;
        let mut hs = HashSet::new();
        if i <= j && k % 2 == 1 {
            hs.insert(v);
        } else {
            hs.insert((v + 1) % 17);
        }
        hs.insert((v + 3) % 17);
        hs.insert((v + 7) % 17);
        hs
    });

    let shape_b = (dj, di);
    let t_b = Array2::<usize>::from_shape_fn(shape_b, |(j, i)| (i * 5 + j * 3 + 11) % 17);

    let result_ndarray = pure_ndarray(&t_a, &t_b);
    let result_tensorism = tensorism(&t_a, &t_b);
    let result_flatten = flatten(&t_a, &t_b);
    assert_eq!(result_ndarray, result_tensorism);
    assert_eq!(result_ndarray, result_flatten);

    c.bench_function("Set belonging - Ndarray", |b| {
        b.iter(|| pure_ndarray(black_box(&t_a), black_box(&t_b)))
    });
    c.bench_function("Set belonging - Tensorism", |b| {
        b.iter(|| tensorism(black_box(&t_a), black_box(&t_b)))
    });
    c.bench_function("Set belonging - Flatten", |b| {
        b.iter(|| flatten(black_box(&t_a), black_box(&t_b)))
    });
    c.bench_function("Set belonging - Flatten accumulator", |b| {
        b.iter(|| flatten_accumulator(black_box(&t_a), black_box(&t_b)))
    });
    c.bench_function("Set belonging - Flatten factorized", |b| {
        b.iter(|| flatten_factorized(black_box(&t_a), black_box(&t_b)))
    });
}

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