use crate::znx_base::{GetZnxBase, ZnxBase, ZnxInfos};
use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxView, alloc_aligned};
use std::marker::PhantomData;

/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
/// stored as a 3D matrix in the DFT domain in a single contiguous array.
/// Each col of the [MatZnxDft] can be seen as a collection of [VecZnxDft].
///
/// [MatZnxDft] is used to permform a vector matrix product between a [VecZnx]/[VecZnxDft] and a [MatZnxDft].
/// See the trait [MatZnxDftOps] for additional information.
pub struct MatZnxDft<D, B> {
    data: D,
    n: usize,
    size: usize,
    rows: usize,
    cols_in: usize,
    cols_out: usize,
    _marker: PhantomData<B>,
}

impl<D, B> ZnxInfos for MatZnxDft<D, B> {
    fn cols(&self) -> usize {
        self.cols_in
    }

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

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

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

    fn sl(&self) -> usize {
        self.n()
    }
}

impl<D, B> DataView for MatZnxDft<D, B> {
    type D = D;
    fn data(&self) -> &Self::D {
        &self.data
    }
}

impl<D, B> DataViewMut for MatZnxDft<D, B> {
    fn data_mut(&mut self) -> &mut Self::D {
        &mut self.data
    }
}

impl<D: AsRef<[u8]>> ZnxView for MatZnxDft<D, FFT64> {
    type Scalar = f64;
}

impl<D, B> MatZnxDft<D, B> {
    pub(crate) fn cols_in(&self) -> usize {
        self.cols_in
    }

    pub(crate) fn cols_out(&self) -> usize {
        self.cols_out
    }
}

impl<D: From<Vec<u8>>, B: Backend> MatZnxDft<D, B> {
    pub(crate) fn bytes_of(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
        unsafe {
            crate::ffi::vmp::bytes_of_vmp_pmat(
                module.ptr,
                (rows * cols_in) as u64,
                (size * cols_out) as u64,
            ) as usize
        }
    }

    pub(crate) fn new(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
        let data: Vec<u8> = alloc_aligned(Self::bytes_of(module, rows, cols_in, cols_out, size));
        Self {
            data: data.into(),
            n: module.n(),
            size,
            rows,
            cols_in,
            cols_out,
            _marker: PhantomData,
        }
    }

    pub(crate) fn new_from_bytes(
        module: &Module<B>,
        rows: usize,
        cols_in: usize,
        cols_out: usize,
        size: usize,
        bytes: impl Into<Vec<u8>>,
    ) -> Self {
        let data: Vec<u8> = bytes.into();
        assert!(data.len() == Self::bytes_of(module, rows, cols_in, cols_out, size));
        Self {
            data: data.into(),
            n: module.n(),
            size,
            rows,
            cols_in,
            cols_out,
            _marker: PhantomData,
        }
    }

    // pub fn from_bytes_borrow(
    //     module: &Module<B>,
    //     rows: usize,
    //     cols_in: usize,
    //     cols_out: usize,
    //     size: usize,
    //     bytes: &mut [u8],
    // ) -> Self {
    //     debug_assert_eq!(
    //         bytes.len(),
    //         Self::bytes_of(module, rows, cols_in, cols_out, size)
    //     );
    //     Self {
    //         inner: ZnxBase::from_bytes_borrow(module.n(), rows, cols_out, size, bytes),
    //         cols_in: cols_in,
    //         cols_out: cols_out,
    //         _marker: PhantomData,
    //     }
    // }
}

impl<D: AsRef<[u8]>> MatZnxDft<D, FFT64> {
    /// Returns a copy of the backend array at index (i, j) of the [MatZnxDft].
    ///
    /// # Arguments
    ///
    /// * `row`: row index (i).
    /// * `col`: col index (j).
    #[allow(dead_code)]
    fn at(&self, row: usize, col: usize) -> Vec<f64> {
        let n: usize = self.n();

        let mut res: Vec<f64> = alloc_aligned(n);

        if n < 8 {
            res.copy_from_slice(&self.raw()[(row + col * self.rows()) * n..(row + col * self.rows()) * (n + 1)]);
        } else {
            (0..n >> 3).for_each(|blk| {
                res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
            });
        }

        res
    }

    #[allow(dead_code)]
    fn at_block(&self, row: usize, col: usize, blk: usize) -> &[f64] {
        let nrows: usize = self.rows();
        let nsize: usize = self.size();
        if col == (nsize - 1) && (nsize & 1 == 1) {
            &self.raw()[blk * nrows * nsize * 8 + col * nrows * 8 + row * 8..]
        } else {
            &self.raw()[blk * nrows * nsize * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
        }
    }
}

pub type MatZnxDftAllocOwned<B> = MatZnxDft<Vec<u8>, B>;