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

/// 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 = "serialization",
    derive(
        mop_common::serde_derive::Serialize,
        mop_common::serde_derive::Deserialize
    )
)]
#[derive(Clone, Debug, Hash, PartialEq, PartialOrd)]
pub struct Css<DS, IS> {
    data: DS,
    dim: Dim<usize>,
    indcs: IS,
}

#[cfg(feature = "std")]
pub type CssVec<D> = Css<Vec<D>, Vec<usize>>;
pub type CssVecArray<DA, IA> = Css<VecArray<DA>, VecArray<IA>>;
pub type CssSlice<'a, D> = Css<&'a [D], &'a [usize]>;
pub type CssSliceMut<'a, D> = Css<&'a mut [D], &'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_vec_array;
    /// assert_eq!(css_vec_array().len(), 10);
    /// ```
    #[inline]
    pub fn len(&self) -> usize {
        self.dim.len()
    }
}

#[cfg(feature = "std")]
impl<D> CssVec<D> {
    pub fn with_vec(len: usize, data: Vec<D>, indcs: Vec<usize>) -> Self {
        Css {
            data,
            dim: len.into(),
            indcs,
        }
    }
}

impl<DA, IA> CssVecArray<DA, IA>
where
    DA: Array,
    IA: Array,
{
    pub fn with_vec_array(len: usize, data: VecArray<DA>, indcs: VecArray<IA>) -> Self {
        Css {
            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)
    }
}

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(feature = "testing")]
impl<T> mop_common::quickcheck::Arbitrary for Css<Vec<T>, Vec<usize>>
where
    T: mop_common::quickcheck::Arbitrary,
{
    #[inline]
    fn arbitrary<G: mop_common::quickcheck::Gen>(g: &mut G) -> Self {
        let size = g.size();
        let len = g.gen_range(0, size);
        let nnz = g.gen_range(0, len + 1);
        Self::new(
            len,
            (0..nnz)
                .map(|_| mop_common::quickcheck::Arbitrary::arbitrary(g))
                .collect(),
            (0..nnz).map(|_| g.gen_range(0, len)).collect(),
        )
    }
}
