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

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

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

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

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

#[cfg(feature = "std")]
impl<T> DrMatrixVec<T> {
  #[cfg(feature = "rand")]
  pub fn new_rnd<F, R>(shape: [usize; 2], rng: &mut R, mut cb: F) -> Self
  where
    F: FnMut(&mut R, [usize; 2]) -> T,
    R: mop_common_deps::rand::Rng,
  {
    let dim: Dim<[usize; 2]> = shape.into();
    let mut data = Vec::with_capacity(dim.rows() * dim.cols());
    for x in 0..dim.rows() {
      for y in 0..dim.cols() {
        data.push(cb(rng, [x, y]));
      }
    }
    DrMatrix { data, dim }
  }

  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,
{
  #[cfg(feature = "rand")]
  pub fn new_rnd<F, R>(shape: [usize; 2], rng: &mut R, mut cb: F) -> Self
  where
    F: FnMut(&mut R, [usize; 2]) -> DA::Item,
    R: mop_common_deps::rand::Rng,
  {
    let dim: Dim<[usize; 2]> = shape.into();
    let data = VecArray::new_rnd(dim.rows() * dim.cols(), rng, |rng, idx| {
      cb(rng, crate::utils::matrix_x_y(dim.cols(), idx))
    });
    DrMatrix { data, dim }
  }

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

impl<DS> DrMatrix<DS>
where
  DS: StDenseStoRef,
{
  /// Creates a new [`DrMatrix`](DrMatrix) from raw parameters.
  ///
  /// # Arguments
  ///
  /// * `shape` - An array containing the number of rows and columns.
  /// * `data` - The matrix data.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use mop_structs::matrix::dr_matrix::DrMatrixArray;
  /// let _ = DrMatrixArray::new([2, 4], [1, 2, 3, 4, 5, 6, 7, 8]);
  /// ```
  ///
  /// # Assertions
  ///
  /// * The length of `data` must be equal the number of rows times the number of columns.
  ///
  /// ```should_panic
  /// use mop_structs::matrix::dr_matrix::DrMatrixRef;
  /// let _ = DrMatrixRef::new([2, 4], &[1, 2, 3]);
  /// ```
  pub fn new(shape: [usize; 2], data: DS) -> Self {
    let dim: Dim<[usize; 2]> = shape.into();
    assert!(dim.rows() * dim.cols() == data.as_slice().len());
    Self::new_unchecked(shape, data)
  }

  /// A faster and unsafe version of [`new`](#method.new).
  pub fn new_unchecked(shape: [usize; 2], data: DS) -> Self {
    DrMatrix {
      data,
      dim: shape.into(),
    }
  }

  /// Converts the inner storage to a generic immutable slice storage.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use mop_structs::{
  ///     doc_tests::dr_matrix_array,
  ///     matrix::dr_matrix::DrMatrixRef
  /// };
  /// assert_eq!(
  ///     dr_matrix_array().as_ref(),
  ///     DrMatrixRef::new(
  ///         [4, 5],
  ///         &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
  ///     )
  /// );
  /// ```
  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 = "rayon")]
  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]
  }
}

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 = "rayon")]
  pub fn row_par_iter_mut(
    &mut self,
  ) -> crate::utils::ParallelIteratorWrapper<DrMatrixRowIterMut<'_, DS::Item>> {
    crate::utils::ParallelIteratorWrapper(self.row_iter_mut())
  }

  pub fn swap(&mut self, a: [usize; 2], b: [usize; 2]) {
    let a_data_idx = self.stride(a[0]) + a[1];
    let b_data_idx = self.stride(b[0]) + b[1];
    self.data_mut().swap(a_data_idx, b_data_idx)
  }

  pub fn two_rows_mut(&mut self, smaller_idx: usize, bigger_idx: usize) -> [&mut [DS::Item]; 2] {
    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(all(feature = "quickcheck", feature = "rand"))]
use mop_common_deps::{
  quickcheck::{Arbitrary, Gen},
  rand::{
    distributions::{Distribution, Standard},
    Rng,
  },
};
#[cfg(all(feature = "quickcheck", feature = "rand"))]
impl<T> Arbitrary for DrMatrixVec<T>
where
  Standard: Distribution<T>,
  T: Arbitrary,
{
  #[inline]
  fn arbitrary<G>(g: &mut G) -> Self
  where
    G: Gen,
  {
    DrMatrixVec::new_rnd(
      [g.gen_range(0, g.size()), g.gen_range(0, g.size())],
      g,
      |g, _| g.gen(),
    )
  }
}

#[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::with_array(DEFAULT_DATA))
  }
}