use std::marker::PhantomData;

use nalgebra::{allocator::Allocator, DefaultAllocator, Dim, SVector};
use rand::{distr::Distribution, Rng, RngCore};
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<D> {
    rng: Box<dyn RngCore>,
    p: f64,
    _phantom: PhantomData<D>
}

impl<D> BinaryStringBitPerturbation<D> {
    pub fn new(p: f64) -> Self {
        Self {
            rng: Box::new(rand::rng()),
            p,
            _phantom: PhantomData
        }
    }
}

impl<D> PerturbationOperator for BinaryStringBitPerturbation<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    type Chromosome = BinaryString<D>;

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

pub struct RandomDistributionPerturbation<const LEN: usize, TDistribution: Distribution<f64>> {
    distribution: TDistribution,
    rng: Box<dyn RngCore>,
    parameter: f64
}

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

    pub fn std_dev(&self) -> f64 {
        self.parameter
    }

    pub fn set_std_dev(&mut self, std_dev: f64) -> Result<f64, NormalError> {
        self.parameter = std_dev;
        self.distribution = Normal::new(0.0, std_dev)?;
        Ok(std_dev)
    }
}

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

    pub fn range(&self) -> f64 {
        self.parameter
    }

    pub fn set_range(&mut self, range: f64) -> Result<f64, uniform::Error> {
        self.parameter = range;
        self.distribution = Uniform::new(-range/2.0, range/2.0)?;
        Ok(range)
    }
}

impl<TDistribution: Distribution<f64>, const LEN: usize> PerturbationOperator for RandomDistributionPerturbation<LEN, 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 struct PatternPerturbation<const LEN: usize> {
    d: f64,
    rng: Box<dyn RngCore>
}

impl<const LEN: usize> PatternPerturbation<LEN> {
    pub fn new(d: f64) -> Self {
        Self {
            d,
            rng: Box::new(rand::rng())
        }
    }
}

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

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        let mut chromosome = chromosome.clone();

        // 1. Choose dimension
        let idx = self.rng.random_range(0..LEN);
        // 2. Direction
        let d = if self.rng.random_bool(0.5) {
            self.d
        } else {
            -self.d
        };

        // Apply
        chromosome[idx] += d;

        chromosome
    }
}

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

    pub fn inner(&self) -> &T {
        &self.perturbation
    }

    pub fn inner_mut(&mut self) -> &mut T {
        &mut self.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))
        }
    }
}

/// Perform given perturbation only with given probability
pub struct MutationPerturbation<T: Clone> {
    perturbation: Box<dyn PerturbationOperator<Chromosome = T>>,
    rng: Box<dyn RngCore>,
    probability: f64
}

impl<T: Clone> MutationPerturbation<T> {
    pub fn new(perturbation: Box<dyn PerturbationOperator<Chromosome = T>>, probability: f64) -> Self {
        Self {
            perturbation,
            rng: Box::new(rand::rng()),
            probability
        }
    }
}

impl<T: Clone> PerturbationOperator for MutationPerturbation<T> {
    type Chromosome = T;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        if self.rng.random_bool(self.probability) {
            self.perturbation.perturb(chromosome)
        } else {
            chromosome.clone()
        }
    }
}

pub struct CombinedPerturbation<T: Clone> {
    perturbations: Vec<Box<dyn PerturbationOperator<Chromosome = T>>>,
}

impl<T: Clone> CombinedPerturbation<T> {
    pub fn new(perturbations: Vec<Box<dyn PerturbationOperator<Chromosome = T>>>) -> Self {
        Self {
            perturbations,
        }
    }
}

impl<T: Clone> PerturbationOperator for CombinedPerturbation<T> {
    type Chromosome = T;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        let mut current_chromosome = chromosome.clone();
        for perturbation in self.perturbations.iter_mut() {
            current_chromosome = perturbation.perturb(&current_chromosome);
        }

        current_chromosome
    }
}

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

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

        assert_eq!(
            *BinaryString::new_dyn(vec![1, 1, 0, 0])
                .perturb(&mut rng, 1.0)
                .vec()
                .iter()
                .map(|&x| x)
                .collect::<Vec<_>>(),
            vec![0, 0, 1, 1]
        );


        assert_eq!(
            *BinaryString::new_dyn(vec![1, 1, 0, 0])
                .perturb(&mut rng, 0.0)
                .vec()
                .iter()
                .map(|&x| x)
                .collect::<Vec<_>>(),
            vec![1, 1, 0, 0]
        );
    }
}