use crate::{
  dim::Dim,
  prelude::{Array, DynDenseStoMut, DynDenseStoRef, StDenseStoMut, StDenseStoRef},
  vec::VecArray,
};

/// Compressed Sparse Storage.
///
/// This struct represents a single sparse row, thtoUses the same logic and storage of a CSR Matrix without rows/pointers.
#[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 Css<DS, IS> {
  data: DS,
  dim: Dim<usize>,
  indcs: IS,
}

pub type CssArray<DA, IA> = Css<DA, IA>;
#[cfg(feature = "std")]
pub type CssVec<T> = Css<Vec<T>, Vec<usize>>;
pub type CssVecArray<DA, IA> = Css<VecArray<DA>, VecArray<IA>>;
pub type CssSlice<'a, T> = Css<&'a [T], &'a [usize]>;
pub type CssSliceMut<'a, T> = Css<&'a mut [T], &'a [usize]>;

impl<DS, IS> Css<DS, IS> {
  pub fn new(len: usize, data: DS, indcs: IS) -> Self {
    Css {
      data,
      dim: len.into(),
      indcs,
    }
  }

  /// The length of the vector.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use mop_structs::doc_tests::css_array;
  /// assert_eq!(css_array().len(), 10);
  /// ```
  #[inline]
  pub fn len(&self) -> usize {
    self.dim.len()
  }
}

#[cfg(feature = "std")]
impl<T> CssVec<T> {
  #[cfg(feature = "rand")]
  pub fn new_rnd<F, R>(len: usize, nnz: usize, rng: &mut R, mut cb: F) -> Self
  where
    F: FnMut(&mut R, usize) -> T,
    R: mop_common_deps::rand::Rng,
  {
    assert!(nnz <= len);
    let mut counter = 0;
    let mut indcs = Vec::with_capacity(nnz);
    while counter < nnz {
      let rnd = rng.gen_range(0, len);
      if indcs.contains(&rnd) == false {
        indcs.push(rnd);
        counter += 1;
      }
    }
    indcs.sort_unstable();
    let data = (0..nnz).map(|idx| cb(rng, indcs[idx])).collect();
    CssVec {
      dim: len.into(),
      data,
      indcs,
    }
  }

  pub fn with_vec(len: usize, data: Vec<T>, indcs: Vec<usize>) -> Self {
    CssVec {
      data,
      dim: len.into(),
      indcs,
    }
  }
}

impl<DA, IA> CssVecArray<DA, IA>
where
  DA: Array,
  IA: Array<Item = usize>,
{
  #[cfg(feature = "rand")]
  pub fn new_rnd<F, R>(len: usize, nnz: usize, rng: &mut R, mut cb: F) -> Self
  where
    F: FnMut(&mut R, usize) -> DA::Item,
    R: mop_common_deps::rand::Rng,
  {
    assert!(nnz <= len);
    let mut counter = 0;
    let mut indcs = VecArray::<IA>::with_capacity();
    while counter < nnz {
      let rnd = rng.gen_range(0, len);
      if indcs.as_slice().contains(&rnd) == false {
        indcs.push(rnd);
        counter += 1;
      }
    }
    indcs.as_mut_slice().sort_unstable();
    let data = VecArray::<DA>::new_rnd(nnz, rng, |rng, idx| cb(rng, indcs.as_slice()[idx]));
    CssVecArray {
      dim: len.into(),
      data,
      indcs,
    }
  }

  pub fn with_vec_array(len: usize, data: VecArray<DA>, indcs: VecArray<IA>) -> Self {
    CssVecArray {
      data,
      dim: len.into(),
      indcs,
    }
  }
}

impl<DS, IS> Css<DS, IS>
where
  DS: StDenseStoRef,
  IS: StDenseStoRef<Item = usize>,
{
  pub fn as_ref(&self) -> CssSlice<'_, DS::Item> {
    Css::new(self.dim.len(), self.data.as_slice(), self.indcs.as_slice())
  }

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

  pub fn get(&self, idx: usize) -> Option<&DS::Item> {
    if let Ok(x) = self.indcs().binary_search(&idx) {
      Some(&self.data()[x])
    } else {
      None
    }
  }

  pub fn indcs(&self) -> &[usize] {
    self.indcs.as_slice()
  }

  pub fn iter(&self) -> impl Iterator<Item = (usize, &DS::Item)> {
    self
      .indcs
      .as_slice()
      .iter()
      .cloned()
      .zip(self.data.as_slice())
  }

  pub fn nnz(&self) -> usize {
    self.data().len()
  }
}

impl<DS, IS> Css<DS, IS>
where
  DS: StDenseStoMut,
  IS: StDenseStoRef<Item = usize>,
{
  pub fn as_mut_slice(&mut self) -> CssSliceMut<'_, DS::Item> {
    Css::new(
      self.dim.len(),
      self.data.as_mut_slice(),
      self.indcs.as_slice(),
    )
  }

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

  pub fn get_mut(&mut self, idx: usize) -> Option<&mut DS::Item> {
    if let Ok(x) = self.indcs().binary_search(&idx) {
      Some(&mut self.data_mut()[x])
    } else {
      None
    }
  }

  pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut DS::Item)> {
    self
      .indcs
      .as_slice()
      .iter()
      .cloned()
      .zip(self.data.as_mut_slice())
  }

  pub fn split_at_mut(
    &mut self,
    idx: usize,
  ) -> (CssSliceMut<'_, DS::Item>, CssSliceMut<'_, DS::Item>) {
    assert!(idx <= self.len());
    let len = self.len();
    let split_idx = self.indcs().binary_search(&idx).unwrap();
    let indcs = self.indcs.as_slice();
    let (first_data, second_data) = self.data.as_mut_slice().split_at_mut(split_idx);
    let first = Css::new(idx, first_data, &indcs[0..split_idx]);
    let second = Css::new(len - idx, second_data, &indcs[split_idx..]);
    (first, second)
  }

  pub fn swap(&mut self, a: usize, b: usize) {
    let a_data_idx = self.indcs().binary_search(&a).unwrap();
    let b_data_idx = self.indcs().binary_search(&b).unwrap();
    self.data_mut().swap(a_data_idx, b_data_idx);
  }
}

impl<DS, IS> Css<DS, IS>
where
  DS: DynDenseStoRef,
{
  pub fn nnz_capacity(&self) -> usize {
    self.data.capacity()
  }
}

impl<DS, IS> Css<DS, IS>
where
  DS: DynDenseStoMut,
  IS: DynDenseStoMut<Item = usize>,
{
  pub fn clear(&mut self) {
    self.data.clear();
    self.indcs.clear();
    *self.dim.len_mut() = 0;
  }

  pub fn extend(&mut self, other: &Self)
  where
    DS::Item: Copy,
  {
    self.data.extend(other.data());
    self.indcs.extend(other.indcs());
    *self.dim.len_mut() += other.len();
  }

  pub fn set(&mut self, idx: usize, data: DS::Item) {
    self.indcs.push(idx);
    self.data.push(data);
    *self.dim.len_mut() += 1;
  }

  pub fn truncate(&mut self, until_idx: usize) {
    self.data.truncate(until_idx);
    self.indcs.truncate(until_idx);
    *self.dim.len_mut() = until_idx;
  }
}

#[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 CssVec<T>
where
  Standard: Distribution<T>,
  T: Arbitrary,
{
  #[inline]
  fn arbitrary<G>(g: &mut G) -> Self
  where
    G: Gen,
  {
    let len = g.gen_range(0, g.size());
    let nnz = g.gen_range(0, len + 1);
    CssVec::new_rnd(len, nnz, g, |g, _| g.gen())
  }
}