use core::{fmt::Debug, ops::RangeInclusive};
use mop_common_defs::TraitCfg;
use mop_common_deps::rand::{
  distributions::{uniform::SampleUniform, Distribution, Uniform},
  Rng,
};
use mop_structs::{
  matrix::csr_matrix::CsrMatrixVecArray,
  matrix::dr_matrix::DrMatrixVecArray,
  prelude::{Matrix, StDenseStoMut},
  utils::matrix_x_y,
  vec::{css::CssVecArray, VecArray},
};

/// A solution domain is a set/tuple of flattened and indexed values.
///
/// A domain creates and modifies solutions.
pub trait SolutionDomain<S>: TraitCfg {
  fn new_rnd_solution<R>(&self, vars_num: RangeInclusive<usize>, rng: &mut R) -> S
  where
    R: Rng;

  fn set_rnd_solution_domain<R>(&self, s: &mut S, idx: usize, rng: &mut R)
  where
    R: Rng;
}

impl<S> SolutionDomain<S> for ()
where
  S: Default,
{
  fn new_rnd_solution<R>(&self, _: RangeInclusive<usize>, _: &mut R) -> S
  where
    R: Rng,
  {
    S::default()
  }

  fn set_rnd_solution_domain<R>(&self, _: &mut S, _: usize, _: &mut R)
  where
    R: Rng,
  {
  }
}

macro_rules! array_impls {
    ($($N:expr),+) => {
        $(

impl<T> SolutionDomain<[T; $N]> for [RangeInclusive<T>; $N]
where
    T: Copy + Debug + SampleUniform + TraitCfg
{
    fn new_rnd_solution<R>(&self, _: RangeInclusive<usize>, rng: &mut R) -> [T; $N]
    where
        R: Rng,
    {
        let mut s: [T; $N] = [
            Uniform::from(*self[0].start()..=*self[0].end()).sample(rng);
            $N
        ];
        for (domain, array_elem) in self.iter().zip(s.iter_mut()).skip(1) {
            *array_elem = Uniform::from(*domain.start()..=*domain.end()).sample(rng);
        }
        s
    }

    fn set_rnd_solution_domain<R>(&self, s: &mut [T; $N], idx: usize, rng: &mut R,)
    where
        R: Rng,
    {
        let domain = &self[idx];
        let domain_value = Uniform::from(*domain.start()..=*domain.end(),).sample(rng);
        s[idx] = domain_value;
    }
}

impl<T> SolutionDomain<CsrMatrixVecArray<[T; $N], [usize; $N]>> for DrMatrixVecArray<[RangeInclusive<T>; $N]>
where
    T: Copy + Debug + SampleUniform + TraitCfg,
{
    fn new_rnd_solution<R>(&self, vars_num: RangeInclusive<usize>, rng: &mut R) -> CsrMatrixVecArray<[T; $N], [usize; $N]>
    where
        R: Rng,
    {
        let nnz = Uniform::from(vars_num).sample(rng);
        CsrMatrixVecArray::new_rnd([self.rows(), self.cols()], nnz, rng, |rng, [x, y]| {
            let range = self.value(x, y);
            Uniform::from(*range.start()..=*range.end()).sample(rng)
        })
    }

    fn set_rnd_solution_domain<R>(&self, s: &mut CsrMatrixVecArray<[T; $N], [usize; $N]>, idx: usize, rng: &mut R,)
    where
        R: Rng,
    {
        let [x, y] = matrix_x_y(self.cols(), idx);
        let range = self.value(x, y);
        *s.value_mut(x, y).unwrap() =
            Uniform::from(*range.start()..=*range.end()).sample(rng);
    }
}

impl<T> SolutionDomain<CssVecArray<[T; $N], [usize; $N]>> for [RangeInclusive<T>; $N]
where
    T: Copy + Debug + SampleUniform + TraitCfg
{
    fn new_rnd_solution<R>(&self, vars_num: RangeInclusive<usize>, rng: &mut R) -> CssVecArray<[T; $N], [usize; $N]>
    where
        R: Rng,
    {
        let nnz = Uniform::from(vars_num).sample(rng);
        CssVecArray::new_rnd($N, nnz, rng, |rng, idx| {
            let range = &self[idx];
            Uniform::from(*range.start()..=*range.end()).sample(rng)
        })
    }

    fn set_rnd_solution_domain<R>(&self, s: &mut CssVecArray<[T; $N], [usize; $N]>, idx: usize, rng: &mut R)
    where
        R: Rng,
    {
        let domain = &self[idx];
        let domain_value = Uniform::from(*domain.start()..=*domain.end()).sample(rng);
        *s.get_mut(idx).unwrap() = domain_value;
    }
}

impl<T> SolutionDomain<DrMatrixVecArray<[T; $N]>> for DrMatrixVecArray<[RangeInclusive<T>; $N]>
where
    T: Copy + Debug + SampleUniform + TraitCfg,
{
    fn new_rnd_solution<R>(&self, _: RangeInclusive<usize>, rng: &mut R) -> DrMatrixVecArray<[T; $N]>
    where
        R: Rng,
    {
        DrMatrixVecArray::new_rnd([self.rows(), self.cols()], rng, |rng, [x, y]| {
            let range = self.value(x, y);
            Uniform::from(*range.start()..=*range.end()).sample(rng)
        })
    }

    fn set_rnd_solution_domain<R>(&self, s: &mut DrMatrixVecArray<[T; $N]>, idx: usize, rng: &mut R)
    where
        R: Rng,
    {
        let [x, y] = matrix_x_y(self.cols(), idx);
        let range = self.value(x, y);
        *s.value_mut(x, y) = Uniform::from(*range.start()..=*range.end()).sample(rng);
    }
}

impl<T> SolutionDomain<VecArray<[T; $N]>> for VecArray<[RangeInclusive<T>; $N]>
where
    T: Copy + Debug + SampleUniform + TraitCfg
{
    fn new_rnd_solution<R>(&self, vars_num: RangeInclusive<usize>, rng: &mut R) -> VecArray<[T; $N]>
    where
        R: Rng,
    {
        let len = Uniform::from(vars_num).sample(rng);
        VecArray::new_rnd(len, rng, |rng, idx| {
            let domain = &self[idx];
            Uniform::from(*domain.start()..=*domain.end()).sample(rng)
        })
    }

    fn set_rnd_solution_domain<R>(&self, s: &mut VecArray<[T; $N]>, idx: usize, rng: &mut R,)
    where
        R: Rng,
    {
        let domain = &self[idx];
        let domain_value = Uniform::from(*domain.start()..=*domain.end()).sample(rng);
        s.as_mut_slice()[idx] = domain_value;
    }
}

        )+
    }
}

array_impls!(
  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26,
  27, 28, 29, 30, 31, 32, 0x40, 0x80, 0x100, 0x200, 0x400, 0x800, 0x1000, 0x2000, 0x4000, 0x8000,
  0x10000
);