use nalgebra::SVector;
use rand::{distr::Distribution, Rng};
use rand_distr::{uniform, Normal, NormalError, Uniform};

use crate::binary_string::BinaryString;

pub trait PerturbationOperator {
    type Chromosome;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome;
}

pub struct BinaryStringBitPerturbation<TRng: Rng> {
    rng: TRng,
    p: f64,
}

impl BinaryStringBitPerturbation<rand::rngs::ThreadRng> {
    pub fn new(p: f64) -> Self {
        Self {
            rng: rand::rng(),
            p
        }
    }
}

impl<TRng: Rng> PerturbationOperator for BinaryStringBitPerturbation<TRng> {
    type Chromosome = BinaryString;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        chromosome.perturb(&mut self.rng, self.p)
    }
}

pub struct RandomDistributionPerturbation<const LEN: usize, TRng: Rng, TDistribution: Distribution<f64>> {
    distribution: TDistribution,
    rng: TRng
}

impl<const LEN: usize> RandomDistributionPerturbation<LEN, rand::rngs::ThreadRng, Normal<f64>> {
    pub fn normal(variance: f64) -> Result<Self, NormalError> {
        Ok(Self {
            distribution: Normal::new(0.0, variance)?,
            rng: rand::rng()
        })
    }
}

impl<const LEN: usize> RandomDistributionPerturbation<LEN, rand::rngs::ThreadRng, Uniform<f64>> {
    pub fn uniform(range: f64) -> Result<Self, uniform::Error> {
        Ok(Self {
            distribution: Uniform::new(-range/2.0, range/2.0)?,
            rng: rand::rng()
        })
    }
}

impl<TRng: Rng, TDistribution: Distribution<f64>, const LEN: usize> PerturbationOperator for RandomDistributionPerturbation<LEN, TRng, TDistribution> {
    type Chromosome = SVector<f64, LEN>;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        chromosome + Self::Chromosome::zeros().map(|_| self.distribution.sample(&mut self.rng))
    }
}

pub enum BoundedPerturbationStrategy {
    /// Trims the value to get a value within bounds
    Trim,
    /// Retries calling the underlying perturbation until
    /// value within bounds is returned. If argument is given,
    /// this is the maximum number of retries to do and then
    /// fall back to trimming. Zero means retry indefinitely.
    Retry(usize)
}

pub struct BoundedPerturbation<const LEN: usize, T: PerturbationOperator<Chromosome = SVector<f64, LEN>>> {
    min_max: SVector<(f64, f64), LEN>,
    strategy: BoundedPerturbationStrategy,
    perturbation: T,
}

impl<const LEN: usize, T: PerturbationOperator<Chromosome = SVector<f64, LEN>>> BoundedPerturbation<LEN, T> {
    pub fn new(
        perturbation: T,
        min: SVector<f64, LEN>,
        max: SVector<f64, LEN>,
        strategy: BoundedPerturbationStrategy
    ) -> Self {
        let min_max = min.zip_map(&max, |min, max| (min, max));
        Self {
            min_max,
            strategy,
            perturbation
        }
    }

    fn within_bounds(&self, chromosome: &SVector<f64, LEN>) -> bool {
        chromosome.iter()
            .zip(self.min_max.iter())
            .all(|(&c, &(min, max))| c <= max && c >= min)
    }

    fn bound(&self, mut chromosome: SVector<f64, LEN>) -> SVector<f64, LEN> {
        chromosome
            .zip_apply(&self.min_max, |c, (min, max)| *c = c.clamp(min, max));

        chromosome
    }

    fn retry_perturb(self: &mut Self, chromosome: &SVector<f64, LEN>, retries: Option<usize>) -> SVector<f64, LEN> {
        let perturbed = self.perturbation.perturb(chromosome);

        if self.within_bounds(&perturbed) {
            return perturbed;
        }

        match retries {
            Some(0) | None => self.bound(perturbed),
            Some(retries) => self.retry_perturb(chromosome, Some(retries - 1))
        }
    }
}

impl<const LEN: usize, T> PerturbationOperator for BoundedPerturbation<LEN, T>
where
    T: PerturbationOperator<Chromosome = SVector<f64, LEN>>
{
    type Chromosome = SVector<f64, LEN>;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        match self.strategy {
            BoundedPerturbationStrategy::Trim => self.retry_perturb(chromosome, None),
            BoundedPerturbationStrategy::Retry(retries) => self.retry_perturb(chromosome, Some(retries))
        }
    }
}

#[cfg(test)]
pub mod tests {
    use crate::binary_string::BinaryString;

    #[test]
    fn test_perturb() {
        let mut rng = rand::rng();

        assert_eq!(
            *BinaryString::new(vec![1, 1, 0, 0])
                .perturb(&mut rng, 1.0)
                .vec(),
            vec![0, 0, 1, 1]
        );

        assert_eq!(
            *BinaryString::new(vec![1, 1, 0, 0])
                .perturb(&mut rng, 0.0)
                .vec(),
            vec![1, 1, 0, 0]
        );
    }
}