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

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<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))
    }
}

#[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]
        );
    }
}