use std::alloc::Layout;
use std::fmt::Display;

use DatumType::F32;
use proptest::arbitrary::Arbitrary;
use proptest::prelude::*;
use proptest::strategy::{BoxedStrategy, Strategy};
use tract_data::internal::*;
use tract_linalg::WeightType;
use tract_linalg::mmm::FusedSpec;
use tract_linalg::mmm::{AsInputValue, EagerPackedInput, MMMInputFormat, MMMInputValue};
use tract_linalg::pack::{PackedFormat, PackingWriter};

proptest::proptest! {
    #[test]
    fn prop(pb in any::<ConvProblem>()) {
        pb.check()
    }
}

#[test]
fn test1() {
    ConvProblem {
        lazy_im2col: false,
        input: tensor3(&[[[1f32]]]),
        filters: tensor4(&[[[[-1f32]]]]),
    }
    .check()
}

#[test]
fn test_axes_0() {
    // CHW HWIO CHW
    // 121 1112 221
    ConvProblem {
        lazy_im2col: false,
        input: tensor3(&[[[0f32], [-1.0]]]),
        filters: tensor4(&[[[[0f32, -1f32]]]]),
    }
    .check()
}

#[test]
fn test_axes_1() {
    ConvProblem {
        lazy_im2col: false,
        input: tensor3(&[[[0f32, 1.]]]),
        filters: tensor4(&[[[[1f32]]]]),
    }
    .check()
}

#[test]
fn test_lazy_0() {
    ConvProblem {
        lazy_im2col: true,
        input: tensor3(&[[[1f32]]]),
        filters: tensor4(&[[[[1f32]]]]),
    }
    .check()
}

#[test]
fn test_lazy_1() {
    ConvProblem {
        lazy_im2col: true,
        input: tensor3(&[[[0f32], [0.], [0.]]]),
        filters: tensor4(&[[[[0f32]]]]),
    }
    .check()
}

#[test]
fn test_lazy_2() {
    ConvProblem {
        lazy_im2col: true,
        input: tensor3(&[[[0f32, 0.], [0., 1.]]]),
        filters: tensor4(&[[[[0f32]], [[1.]]]]),
    }
    .check()
}

#[test]
fn test_lazy_3() {
    // CHW HWIO CHW
    // 212 1221 111
    // im2col: k=4, n=1, k <- kh, kw, c
    // 0 X X X X kh=0, kw=0, c=0
    // 1 X X X X kh=0, kw=0, c=1
    // 0 X X X X kh=0, kw=1, c=0
    // 0 X X X X kh=0, kw=1, c=1
    ConvProblem {
        lazy_im2col: true,
        input: tensor3(&[[[0f32, 0.]], [[1., 0.]]]),
        filters: tensor4(&[[[[0f32], [0.]], [[1.], [0.]]]]),
    }
    .check()
}

#[test]
fn test_eager_asan_0() {
    ConvProblem {
        lazy_im2col: false,
        input: tensor(vec![3, 3, 5]),
        filters: tensor(vec![3, 3, 3, 1]),
    }
    .check()
}

// 2D valid, no group, no dil, no stride, HWIO, CHW
#[derive(Clone, Debug)]
pub struct ConvProblem {
    pub lazy_im2col: bool,
    pub input: Tensor,
    pub filters: Tensor,
}

fn mknhw(filters: &[usize], input: &[usize]) -> (usize, usize, usize, usize, usize) {
    let m = filters[3];
    let k = filters[0..3].iter().product::<usize>();
    let h = input[1] - filters[0] + 1;
    let w = input[2] - filters[1] + 1;
    let n = h * w;
    (m, k, n, h, w)
}

impl ConvProblem {
    fn reference(&self) -> Tensor {
        let (m, _, _, h, w) = mknhw(self.filters.shape(), self.input.shape());
        let output_shape = [m, h, w];
        let mut output = Tensor::zero::<f32>(&output_shape).unwrap();
        let mut output_plain = output.try_as_plain_mut().unwrap();
        let mut output_view = output_plain.to_array_view_mut::<f32>().unwrap();
        let input_view = self.input.to_plain_array_view::<f32>().unwrap();
        let filters_view = self.filters.to_plain_array_view::<f32>().unwrap();
        for geo_out in tract_ndarray::indices(&output_shape[1..]) {
            for ker_geo in tract_ndarray::indices(&self.filters.shape()[0..2]) {
                for ci in 0..self.filters.shape()[2] {
                    for co in 0..self.filters.shape()[3] {
                        let output_coord = [co, geo_out[0], geo_out[1]];
                        let input_coord = [ci, geo_out[0] + ker_geo[0], geo_out[1] + ker_geo[1]];
                        let ker_coord = [ker_geo[0], ker_geo[1], ci, co];
                        output_view[output_coord] +=
                            filters_view[ker_coord] * input_view[input_coord];
                    }
                }
            }
        }
        output
    }

    pub fn tract(&self) -> TractResult<Tensor> {
        let (m, k, n, h, w) = mknhw(self.filters.shape(), self.input.shape());
        let output_shape = [m, h, w];
        let internal_output_shape = [m, h * w];
        let mmm = tract_linalg::ops()
            .mmm(F32, Some(m), Some(k), Some(n))
            .unwrap();
        let output = Tensor::zero::<f32>(&internal_output_shape)?;
        let reshaped_filters = self.filters.clone().into_shape(&[k, m])?;
        let (a_pack, b_pack) = &mmm.packings()[0];
        let a = a_pack.prepare_one(&reshaped_filters, 0, 1)?;
        unsafe {
            let im2col: Box<dyn MMMInputValue> = if self.lazy_im2col {
                LazyIm2colSpec {
                    full_kernel_shape: self.filters.shape().into(),
                    packer: b_pack.downcast_ref::<PackedFormat>().unwrap().clone(),
                }
                .wrap(&self.input.view())
            } else {
                EagerIm2colSpec {
                    full_kernel_shape: self.filters.shape().into(),
                    packer: b_pack.downcast_ref::<PackedFormat>().unwrap().clone(),
                }
                .wrap(&self.input.view())
            };
            let c_store = mmm.c_view(Some(0), Some(1)).wrap(&output.view());
            mmm.run(
                m,
                n,
                &[
                    FusedSpec::AddMatMul {
                        a: AsInputValue::Owned(a),
                        b: AsInputValue::Owned(im2col),
                        packing: 0,
                    },
                    FusedSpec::Store(c_store),
                ],
            )
            .unwrap()
        }
        output.into_shape(&output_shape)
    }

    fn check(&self) {
        let expected = self.reference();
        let found = self.tract().unwrap();
        if found.close_enough(&expected, true).is_err() {
            println!("found: ");
            println!("{:?}", found.to_plain_array_view::<f32>().unwrap());
            println!("expected: ");
            println!("{:?}", expected.to_plain_array_view::<f32>().unwrap());
        }
        found.close_enough(&expected, true).unwrap()
    }
}

impl Arbitrary for ConvProblem {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;
    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
        (
            any::<bool>(),
            1..4usize,
            1..4usize,
            1..4usize,
            1..4usize,
            0..3usize,
            0..3usize,
        )
            .prop_map(|(eager_im2col, h, w, i, o, extra_h, extra_w)| {
                let filters = tensor(vec![h, w, i, o]);
                let input = tensor(vec![i, h + extra_h, w + extra_w]);
                ConvProblem {
                    lazy_im2col: eager_im2col,
                    filters,
                    input,
                }
            })
            .boxed()
    }
}

fn tensor(shape: Vec<usize>) -> Tensor {
    let mut tensor = Tensor::zero::<f32>(&shape).unwrap();
    tensor
        .try_as_plain_mut()
        .unwrap()
        .as_slice_mut::<f32>()
        .unwrap()
        .iter_mut()
        .enumerate()
        .for_each(|(ix, x)| *x = ix as f32);
    tensor
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
struct EagerIm2colSpec {
    packer: PackedFormat,
    full_kernel_shape: TVec<usize>,
}

impl EagerIm2colSpec {
    fn wrap(&self, input: &TensorView) -> Box<dyn MMMInputValue> {
        let (_, k, n, h, w) = mknhw(&self.full_kernel_shape, input.shape());
        // let input = input.to_array_view::<f32>().unwrap();
        let ci = input.shape()[0];
        let kh = self.full_kernel_shape[0];
        let kw = self.full_kernel_shape[1];
        let im2col = tract_ndarray::Array5::<f32>::from_shape_fn(
            [kh, kw, ci, h, w],
            |(kh, kw, ci, h, w)| *input.at([ci, h + kh, w + kw]).unwrap(),
        )
        .into_shape_with_order([k, n])
        .unwrap();
        Box::new(EagerIm2col {
            im2col: im2col.into_tensor(),
            packer: self.packer.clone(),
            k,
        })
    }
}

impl Display for EagerIm2colSpec {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "EagerIm2colSpec")
    }
}

impl MMMInputFormat for EagerIm2colSpec {
    fn prepare_tensor(&self, _t: &Tensor, _k_axis: usize, _mn_axis: usize) -> TractResult<Tensor> {
        todo!();
    }

    fn precursor(&self) -> WeightType {
        WeightType::Plain(f32::datum_type())
    }

    fn k_alignment(&self) -> usize {
        1
    }

    fn r(&self) -> usize {
        self.packer.r()
    }

    fn mem_size(&self, _k: TDim, _mn: TDim) -> TDim {
        unimplemented!()
    }

    fn extract_at_mn_f16(
        &self,
        _data: &EagerPackedInput,
        _mn: usize,
        _slice: &mut [f16],
    ) -> TractResult<()> {
        todo!();
    }

    fn extract_at_mn_f32(
        &self,
        _data: &EagerPackedInput,
        _mn: usize,
        _slice: &mut [f32],
    ) -> TractResult<()> {
        todo!();
    }

    fn prepare_one(
        &self,
        _t: &Tensor,
        _k_axis: usize,
        _mn_axis: usize,
    ) -> TractResult<Box<dyn MMMInputValue>> {
        todo!()
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
struct EagerIm2col {
    packer: PackedFormat,
    im2col: Tensor,
    k: usize,
}

impl Display for EagerIm2col {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "eager")
    }
}

impl MMMInputValue for EagerIm2col {
    fn scratch_panel_buffer_layout(&self) -> Option<std::alloc::Layout> {
        Some(
            Layout::from_size_align(
                self.packer.single_panel_len(self.k) * f32::datum_type().size_of(),
                self.packer.alignment(),
            )
            .unwrap(),
        )
    }

    fn panel_bytes(&self, i: usize, buffer: Option<*mut u8>) -> TractResult<*const u8> {
        let buffer = buffer.unwrap();
        let mn = self.im2col.shape()[1];
        unsafe {
            self.packer.pack_t::<f32>(
                buffer as _,
                self.im2col.as_ptr().unwrap(),
                mn,
                mn as isize,
                1,
                0..self.k,
                (i * self.packer.r)..((i + 1) * self.packer.r),
            );
        }
        Ok(buffer)
    }

    fn k(&self) -> usize {
        self.k
    }

    fn mn(&self) -> usize {
        self.im2col.shape()[1]
    }

    fn format(&self) -> &dyn tract_linalg::mmm::MMMInputFormat {
        &self.packer
    }

    fn exotic_fact(&self) -> &dyn ExoticFact {
        unimplemented!()
    }

    fn extract_at_mn_f16(&self, _mn: usize, _slice: &mut [f16]) -> TractResult<()> {
        unimplemented!()
    }

    fn extract_at_mn_f32(&self, _mn: usize, _slice: &mut [f32]) -> TractResult<()> {
        unimplemented!()
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
struct LazyIm2colSpec {
    packer: PackedFormat,
    full_kernel_shape: TVec<usize>,
}

impl LazyIm2colSpec {
    fn wrap(&self, input: &TensorView) -> Box<dyn MMMInputValue> {
        let (_, _, _, h, w) = mknhw(&self.full_kernel_shape, input.shape());
        let kh = self.full_kernel_shape[0];
        let kw = self.full_kernel_shape[1];
        let ci = self.full_kernel_shape[2];
        let input_strides = input.strides();
        let k_offsets = (0..kh as isize)
            .flat_map(|kh| {
                (0..kw as isize).flat_map(move |kw| {
                    (0..ci as isize).map(move |ci| {
                        ci * input_strides[0] + kh * input_strides[1] + kw * input_strides[2]
                    })
                })
            })
            .collect();
        let n_offsets = (0..h as isize)
            .flat_map(|h| (0..w as isize).map(move |w| h * input_strides[1] + w * input_strides[2]))
            .collect();
        unsafe {
            Box::new(LazyIm2col {
                spec: self.clone(),
                image: input.as_ptr_unchecked(),
                k_offsets,
                n_offsets,
                packer: self.packer.clone(),
            })
        }
    }
}

impl Display for LazyIm2colSpec {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "LazyIm2colSpec")
    }
}

impl MMMInputFormat for LazyIm2colSpec {
    fn prepare_tensor(&self, _t: &Tensor, _k_axis: usize, _mn_axis: usize) -> TractResult<Tensor> {
        todo!();
    }
    fn prepare_one(
        &self,
        _t: &Tensor,
        _k_axis: usize,
        _mn_axis: usize,
    ) -> TractResult<Box<dyn MMMInputValue>> {
        todo!();
    }

    fn precursor(&self) -> WeightType {
        WeightType::Plain(f32::datum_type())
    }

    fn k_alignment(&self) -> usize {
        1
    }

    fn r(&self) -> usize {
        self.packer.r()
    }

    fn mem_size(&self, _k: TDim, _mn: TDim) -> TDim {
        unimplemented!()
    }

    fn extract_at_mn_f16(
        &self,
        _data: &EagerPackedInput,
        _mn: usize,
        _slice: &mut [f16],
    ) -> TractResult<()> {
        todo!();
    }

    fn extract_at_mn_f32(
        &self,
        _data: &EagerPackedInput,
        _mn: usize,
        _slice: &mut [f32],
    ) -> TractResult<()> {
        todo!();
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
struct LazyIm2col {
    spec: LazyIm2colSpec,
    packer: PackedFormat,
    image: *const f32,
    n_offsets: Vec<isize>,
    k_offsets: Vec<isize>,
}
unsafe impl Send for LazyIm2col {}
unsafe impl Sync for LazyIm2col {}

impl Display for LazyIm2col {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "lazy")
    }
}

impl MMMInputValue for LazyIm2col {
    fn scratch_panel_buffer_layout(&self) -> Option<std::alloc::Layout> {
        Some(
            Layout::from_size_align(
                self.packer
                    .single_panel_len(self.k_offsets.len() * f32::datum_type().size_of()),
                self.packer.alignment(),
            )
            .unwrap(),
        )
    }

    fn panel_bytes(&self, i: usize, buffer: Option<*mut u8>) -> TractResult<*const u8> {
        let buffer = buffer.unwrap() as *mut f32;
        let mn_end = ((i + 1) * self.packer.r).min(self.n_offsets.len());
        let n_range = (i * self.packer.r)..mn_end;
        let k = self.k_offsets.len();
        unsafe {
            let mut writer = self.packer.write_with_k_outer(buffer, k, n_range.len());
            for k in 0..k {
                for n in n_range.clone() {
                    writer.write(
                        *self.image.offset(
                            self.n_offsets.get_unchecked(n) + self.k_offsets.get_unchecked(k),
                        ),
                    )
                }
            }
        }
        Ok(buffer as _)
    }

    fn k(&self) -> usize {
        self.k_offsets.len()
    }

    fn mn(&self) -> usize {
        self.n_offsets.len()
    }

    fn format(&self) -> &dyn MMMInputFormat {
        &self.spec
    }

    fn exotic_fact(&self) -> &dyn ExoticFact {
        unimplemented!()
    }

    fn extract_at_mn_f16(&self, _mn: usize, _slice: &mut [f16]) -> TractResult<()> {
        unimplemented!()
    }

    fn extract_at_mn_f32(&self, _mn: usize, _slice: &mut [f32]) -> TractResult<()> {
        unimplemented!()
    }
}
