use std::{any::Any, borrow::{Borrow, BorrowMut}, 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 AnyPerturbationOperator: Any {
    fn as_any(&self) -> &dyn Any;
    fn as_any_mut(&mut self) -> &mut dyn Any;
}

pub enum Wrapped<'a, T> {
    Single,
    Wrapped(&'a dyn PerturbationOperator<Chromosome = T>),
    ListWrapped(Vec<&'a dyn PerturbationOperator<Chromosome = T>>),
}

pub enum WrappedMut<'a, T> {
    Single,
    Wrapped(&'a mut dyn PerturbationOperator<Chromosome = T>),
    ListWrapped(Vec<&'a mut dyn PerturbationOperator<Chromosome = T>>),
}

pub trait PerturbationOperator: AnyPerturbationOperator {
    type Chromosome;

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore);

    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
        Wrapped::Single
    }

    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
        WrappedMut::Single
    }
}

impl<T: PerturbationOperator> AnyPerturbationOperator for T
where
    T: Any + 'static
{
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}

pub struct IdentityPerturbation<TChromosome> {
    _phantom: PhantomData<TChromosome>
}

impl<TChromosome: 'static> PerturbationOperator for IdentityPerturbation<TChromosome> {
    type Chromosome = TChromosome;

    fn perturb(&self, _: &mut Self::Chromosome, _: &mut dyn RngCore) {
        // Do nothing.
    }
}

pub struct BinaryStringBitPerturbation<D> {
    p: f64,
    _phantom: PhantomData<D>
}

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

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

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        chromosome.perturb(rng, self.p);
    }
}

pub struct BinaryStringSingleBitPerturbation<D> {
    _phantom: PhantomData<D>
}

impl<D> BinaryStringSingleBitPerturbation<D> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData
        }
    }
}

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

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        let bit_range = 0..chromosome.vec.len();
        let flip_bit = rng.random_range(bit_range);

        chromosome.vec[flip_bit] = 1 - chromosome.vec[flip_bit];
    }
}

pub struct BinaryStringFlipPerturbation<D> {
    _phantom: PhantomData<D>
}

impl<D> BinaryStringFlipPerturbation<D> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData
        }
    }
}

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

    fn perturb(&self, chromosome: &mut Self::Chromosome, _: &mut dyn RngCore) {
        chromosome.vec
            .apply(|c| *c = 1 - *c);
    }
}

pub struct RandomDistributionPerturbation<const LEN: usize, TDistribution: Distribution<f64>> {
    distribution: TDistribution,
    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)?,
            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)?,
            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> + 'static, const LEN: usize> PerturbationOperator for RandomDistributionPerturbation<LEN, TDistribution> {
    type Chromosome = SVector<f64, LEN>;

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        *chromosome += Self::Chromosome::zeros().map(|_| self.distribution.sample(rng));
    }
}

pub struct PatternPerturbation<const LEN: usize> {
    d: f64
}

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

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

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        // 1. Choose dimension
        let idx = rng.random_range(0..LEN);
        // 2. Direction
        let d = if rng.random_bool(0.5) {
            self.d
        } else {
            -self.d
        };

        // Apply
        chromosome[idx] += d;
    }
}

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, chromosome: &mut SVector<f64, LEN>, retries: Option<usize>, rng: &mut dyn RngCore) {
        let mut perturbed = chromosome.clone();
        self.perturbation.perturb(&mut perturbed, rng);

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

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

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, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        match self.strategy {
            BoundedPerturbationStrategy::Trim => self.retry_perturb(chromosome, None, rng),
            BoundedPerturbationStrategy::Retry(retries) => self.retry_perturb(chromosome, Some(retries), rng)
        }
    }

    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
        Wrapped::Wrapped(&self.perturbation)
    }

    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
        WrappedMut::Wrapped(&mut self.perturbation)
    }
}

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

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

    pub fn apply_to_mutations(base_perturbation: &mut dyn PerturbationOperator<Chromosome = T>, apply: &mut dyn FnMut(&mut Self)) {
        if let Some(mutation) = base_perturbation.as_any_mut().downcast_mut::<MutationPerturbation<T>>() {
            apply(mutation);
        }

        match base_perturbation.wrapped_mut() {
            WrappedMut::Single => (),
            WrappedMut::Wrapped(wrapped) => {
                Self::apply_to_mutations(wrapped, apply);
            },
            WrappedMut::ListWrapped(wrapped) => {
                for wrapped in wrapped {
                    Self::apply_to_mutations(wrapped, apply);
                }
            }
        };
    }
}

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

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        if rng.random_bool(self.probability) {
            self.perturbation.perturb(chromosome, rng);
        }
    }

    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
        Wrapped::Wrapped(self.perturbation.as_ref())
    }

    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
        WrappedMut::Wrapped(self.perturbation.as_mut())
    }
}

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

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

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

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        for perturbation in self.perturbations.iter() {
            perturbation.perturb(chromosome, rng);
        }
    }

    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
        Wrapped::ListWrapped(
            self.perturbations
                .iter()
                .map(|p| p.as_ref()).collect())
    }

    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
        WrappedMut::ListWrapped(
            self.perturbations
                .iter_mut()
                .map(|p| p.as_mut()).collect())
    }
}

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

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

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

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