mod dr_matrix_row_constructor;
mod dr_matrix_row_iter_impls;
#[cfg(feature = "parallelization")]
mod dr_matrix_row_par_iter_impls;

pub use self::{
    dr_matrix_row_constructor::DrMatrixRowConstructor,
    dr_matrix_row_iter_impls::{DrMatrixRowIter, DrMatrixRowIterMut},
};
use crate::{
    prelude::{Array, DynDenseStoMut, DynDenseStoRef, Matrix, StDenseStoMut, StDenseStoRef},
    vec::VecArray,
    Dim,
};

#[cfg_attr(
    feature = "serialization",
    derive(
        mop_common::serde_derive::Serialize,
        mop_common::serde_derive::Deserialize
    )
)]
#[derive(Clone, Debug, Hash, PartialEq, PartialOrd)]
pub struct DrMatrix<DS> {
    data: DS,
    dim: Dim<[usize; 2]>,
}

pub type DrMatrixMut<'a, D> = DrMatrix<&'a mut [D]>;
pub type DrMatrixRef<'a, D> = DrMatrix<&'a [D]>;
#[cfg(feature = "std")]
pub type DrMatrixVec<D> = DrMatrix<Vec<D>>;
pub type DrMatrixVecArray<DA> = DrMatrix<VecArray<DA>>;

impl<DS> DrMatrix<DS> {
    pub fn new(shape: [usize; 2], data: DS) -> Self {
        DrMatrix {
            data,
            dim: shape.into(),
        }
    }
}

#[cfg(feature = "std")]
impl<D> DrMatrixVec<D> {
    pub fn with_vec_capacity(shape: [usize; 2]) -> Self {
        DrMatrix {
            data: Vec::with_capacity(shape[0] * shape[1]),
            dim: [0, shape[1]].into(),
        }
    }
}

impl<DA> DrMatrixVecArray<DA>
where
    DA: Array,
{
    pub fn with_vec_array(cols: usize) -> Self {
        DrMatrix {
            data: VecArray::with_capacity(),
            dim: [0, cols].into(),
        }
    }
}

impl<DS> DrMatrix<DS>
where
    DS: StDenseStoRef,
{
    pub fn as_ref(&self) -> DrMatrixRef<'_, DS::Item> {
        DrMatrix::new([self.dim.rows(), self.dim.cols()], self.data.as_slice())
    }

    pub fn data(&self) -> &[DS::Item] {
        &self.data.as_slice()
    }

    pub fn row(&self, idx: usize) -> &[DS::Item] {
        let stride = self.stride(idx);
        &self.data()[stride..stride + self.cols()]
    }

    pub fn row_iter(&self) -> DrMatrixRowIter<'_, DS::Item> {
        DrMatrixRowIter::new([self.rows(), self.cols()], self.data().as_ptr())
    }

    #[cfg(feature = "parallelization")]
    pub fn row_par_iter(
        &self,
    ) -> crate::utils::ParallelIteratorWrapper<DrMatrixRowIter<'_, DS::Item>> {
        crate::utils::ParallelIteratorWrapper(self.row_iter())
    }

    pub fn value(&self, row_idx: usize, col_idx: usize) -> &DS::Item {
        &self.data()[self.stride(row_idx) + col_idx]
    }

    #[inline]
    fn stride(&self, idx: usize) -> usize {
        self.cols() * idx
    }
}

impl<DS> DrMatrix<DS>
where
    DS: StDenseStoMut,
{
    pub fn data_mut(&mut self) -> &mut [DS::Item] {
        self.data.as_mut_slice()
    }

    pub fn row_mut(&mut self, idx: usize) -> &mut [DS::Item] {
        let stride = self.stride(idx);
        &mut self.data.as_mut_slice()[stride..stride + self.dim.cols()]
    }

    pub fn row_iter_mut(&mut self) -> DrMatrixRowIterMut<'_, DS::Item> {
        DrMatrixRowIterMut::new([self.rows(), self.cols()], self.data_mut().as_mut_ptr())
    }

    #[cfg(feature = "parallelization")]
    pub fn row_par_iter_mut(
        &mut self,
    ) -> crate::utils::ParallelIteratorWrapper<DrMatrixRowIterMut<'_, DS::Item>> {
        crate::utils::ParallelIteratorWrapper(self.row_iter_mut())
    }

    pub fn two_rows_mut(
        &mut self,
        smaller_idx: usize,
        bigger_idx: usize,
    ) -> (&mut [DS::Item], &mut [DS::Item]) {
        let bigger = self.stride(bigger_idx);
        let smaller = self.stride(smaller_idx);
        let (first, second) = self.data.as_mut_slice().split_at_mut(bigger);
        (
            &mut first[smaller..smaller + self.dim.cols()],
            &mut second[0..self.dim.cols()],
        )
    }

    pub fn value_mut(&mut self, row_idx: usize, col_idx: usize) -> &mut DS::Item {
        let stride = self.stride(row_idx);
        &mut self.data_mut()[stride + col_idx]
    }
}

impl<DS> DrMatrix<DS>
where
    DS: DynDenseStoRef,
{
    #[inline]
    pub fn rows_capacity(&self) -> usize {
        self.data.capacity() / self.dim.cols()
    }
}

impl<DS> DrMatrix<DS>
where
    DS: DynDenseStoMut,
{
    pub fn clear(&mut self) {
        self.data.clear();
        *self.dim.rows_mut() = 0;
    }

    pub fn extend<AD>(&mut self, other: &DrMatrix<AD>)
    where
        AD: DynDenseStoMut<Item = DS::Item>,
        DS::Item: Copy,
    {
        assert!(self.cols() == other.cols());
        self.data.extend(other.data());
        *self.dim.rows_mut() += other.rows();
    }

    pub fn extend_from_clone<AD>(&mut self, other: &DrMatrix<AD>)
    where
        AD: DynDenseStoMut<Item = DS::Item>,
        DS::Item: Clone,
    {
        assert!(self.cols() == other.cols());
        self.data.extend_from_clone(other.data());
        *self.dim.rows_mut() += other.rows();
    }

    pub fn row_constructor(&mut self) -> DrMatrixRowConstructor<'_, DS> {
        DrMatrixRowConstructor::new(&mut self.data, &mut self.dim)
    }

    pub fn truncate(&mut self, until_row_idx: usize) {
        self.data.truncate(self.dim.cols() * until_row_idx);
        *self.dim.rows_mut() = until_row_idx;
    }
}

impl<DS> Matrix for DrMatrix<DS> {
    #[inline]
    fn cols(&self) -> usize {
        self.dim.cols()
    }

    #[inline]
    fn rows(&self) -> usize {
        self.dim.rows()
    }
}

#[cfg(feature = "testing")]
impl<T> mop_common::quickcheck::Arbitrary for DrMatrix<Vec<T>>
where
    T: mop_common::quickcheck::Arbitrary,
{
    #[inline]
    fn arbitrary<G: mop_common::quickcheck::Gen>(g: &mut G) -> Self {
        let size = g.size();
        let cols = g.gen_range(0, size);
        let rows = g.gen_range(0, size);

        Self::new(
            [rows, cols],
            (0..rows * cols)
                .map(|_| mop_common::quickcheck::Arbitrary::arbitrary(g))
                .collect(),
        )
    }
}

#[cfg(test)]
mod tests {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    const DEFAULT_DATA: [i32; 20] = [
        1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
        11, 12, 13, 14, 15, 16, 17, 18, 19, 20
    ];
    const DEFAULT_SHAPE: [usize; 2] = [4, 5];

    use crate::{matrix::dr_matrix::DrMatrix, prelude::*, vec::VecArray};

    #[test]
    pub fn dr_matrix_alloc_clear() {
        let mut dcca = default_dr_matrix_alloc();
        dcca.clear();
        assert_eq!(dcca.cols(), DEFAULT_SHAPE[1]);
        assert_eq!(dcca.data(), &[]);
        assert_eq!(dcca.rows(), 0);
    }

    #[test]
    pub fn dr_matrix_alloc_cols() {
        assert_eq!(default_dr_matrix_alloc().cols(), DEFAULT_SHAPE[1]);
    }

    #[test]
    pub fn dr_matrix_alloc_data() {
        assert_eq!(default_dr_matrix_alloc().data(), &DEFAULT_DATA);
    }

    #[test]
    pub fn dr_matrix_alloc_data_mut() {
        assert_eq!(default_dr_matrix_alloc().data_mut(), &mut DEFAULT_DATA);
    }

    #[test]
    pub fn dr_matrix_alloc_row() {
        let ddma = default_dr_matrix_alloc();
        for row_idx in 0..4 {
            let starting_row_value = row_idx * 5 + 1;
            assert_eq!(
                ddma.row(row_idx as usize),
                &[
                    starting_row_value,
                    starting_row_value + 1,
                    starting_row_value + 2,
                    starting_row_value + 3,
                    starting_row_value + 4
                ]
            );
        }
    }

    #[test]
    pub fn dr_matrix_alloc_rows() {
        assert_eq!(default_dr_matrix_alloc().rows(), DEFAULT_SHAPE[0]);
    }

    #[test]
    pub fn dr_matrix_alloc_truncate() {
        let mut ddma = default_dr_matrix_alloc();
        ddma.truncate(2);
        assert_eq!(ddma.cols(), DEFAULT_SHAPE[1]);
        assert_eq!(ddma.data(), &DEFAULT_DATA[0..10]);
        assert_eq!(ddma.rows(), 2);
    }

    #[test]
    pub fn dr_matrix_alloc_value() {
        let ddma = default_dr_matrix_alloc();
        for row_idx in 0..4 {
            let starting_row_value = row_idx * 5 + 1;
            for col_idx in 0..5 {
                assert_eq!(
                    *ddma.value(row_idx as usize, col_idx as usize),
                    starting_row_value + col_idx
                );
            }
        }
    }

    fn default_dr_matrix_alloc() -> DrMatrix<VecArray<[i32; 20]>> {
        DrMatrix::new([4, 5], VecArray::new(DEFAULT_DATA))
    }
}
