use std::{collections::VecDeque, convert::Infallible, error::Error};

use thiserror::Error;

use crate::{comparison::BetterThanOperator, crossover::Crossover, evolution::EvolutionStats, fitness::FitnessFunction, pairing::Pairing, perturbation::PerturbationOperator, population::EvaluatedPopulation, replacement::Replacement, selection::Selection};

pub trait ConstraintFunction {
    type Chromosome;
    type Out;
    type Err: Error + 'static;

    fn evaluate(&self, chromosome: &Self::Chromosome) -> Result<Self::Out, Self::Err>;
    fn is_feasible(&self, chromosome: &Self::Chromosome) -> Result<bool, Self::Err>;
}

pub struct LowerThanConstraintFunction<TChromosome, TOut> {
    fun: Box<dyn Fn(&TChromosome) -> TOut>
}

impl<TChromosome, TOut> LowerThanConstraintFunction<TChromosome, TOut> {
    pub fn new(fun: Box<dyn Fn(&TChromosome) -> TOut>) -> Self {
        Self {
            fun
        }
    }
}

impl<TChromosome, TOut: Default + PartialOrd> ConstraintFunction for LowerThanConstraintFunction<TChromosome, TOut> {
    type Chromosome = TChromosome;
    type Out = TOut;
    type Err = Infallible;

    fn evaluate(&self, chromosome: &Self::Chromosome) -> Result<Self::Out, Self::Err> {
        Ok((self.fun)(chromosome))
    }

    fn is_feasible(&self, chromosome: &Self::Chromosome) -> Result<bool, Self::Err> {
        Ok(self.evaluate(chromosome)? <= Default::default())
    }
}

pub struct ConstrainedFitnessFunction<'a,
    const CONSTRAINTS: usize,
    TIn,
    TOut,
    TFitness: FitnessFunction<In = TIn, Out = TOut>,
    TConstraint: ConstraintFunction<Chromosome = TIn, Out = TOut>> {
    fitness: &'a TFitness,
    constraints: [&'a TConstraint; CONSTRAINTS],
    constraint_weights: Vec<TOut>
}

#[derive(Error, Debug)]
pub enum ConstrainedFitnessErr<T: Error, U: Error> {
    #[error("An error that came from fitness function")]
    FitnessErr(T),
    #[error("An error that came from constraint function")]
    ConstraintErr(U)
}

impl <'a,
      const CONSTRAINTS: usize,
      TOut: std::ops::Mul<Output = TOut> + std::ops::AddAssign + Copy,
      TIn,
      TFitness: FitnessFunction<In = TIn, Out = TOut>,
      TConstraint: ConstraintFunction<Chromosome = TIn, Out = TOut>>
    FitnessFunction for ConstrainedFitnessFunction<'a, CONSTRAINTS, TIn, TOut, TFitness, TConstraint> {
    type In = TFitness::In;
    type Out = TOut;
    type Err = ConstrainedFitnessErr<TFitness::Err, TConstraint::Err>;

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        let mut fit = match self.fitness.fit(inp) {
            Ok(fit) => fit,
            Err(err) =>
                return Err(ConstrainedFitnessErr::FitnessErr(err))
        };

        for (constraint, weight) in self.constraints.iter().zip(self.constraint_weights.iter()) {
            fit += weight.clone() * match constraint.evaluate(inp) {
                Ok(constraint) => constraint,
                Err(err) =>
                    return Err(ConstrainedFitnessErr::ConstraintErr(err))
            };
        }

        Ok(fit)
    }
}

// TODO: currently these functions do recalculate the constraints for each chromosome.
// This is suboptimal. It could be solved by changing the result of fitness function to
// a tuple, where the second element of the tuple would be evaluations of the constraints.
// For this case, it would be the best if the number of constraints has been determined
// by a generic. Then, no dynamic allocation is necessary for each element.

pub fn evolve_constraint_penalty_weight_k
    <TChromosome: Clone,
     const CONSTRAINTS: usize,
     const DParents: usize,
     TSelection: Selection<TChromosome, f64>,
     TFitness: FitnessFunction<In = TChromosome, Out = f64>,
     TConstraint: ConstraintFunction<Chromosome = TChromosome, Out = f64>,
     TPairing: Pairing<DParents, Chromosome = TChromosome, Out = f64>,
     TCrossover: Crossover<DParents, Chromosome = TChromosome, Out = f64>,
     TReplacement: Replacement<TChromosome, f64>,
     TPerturbation: PerturbationOperator<Chromosome = TChromosome>>(
        k: usize,
        n: usize,
        beta_1: f64,
        beta_2: f64,
        better_than: &impl BetterThanOperator<f64>
    ) -> impl FnMut(
        usize,
        &EvolutionStats<TChromosome, f64>,
        &EvaluatedPopulation<TChromosome, f64>,

        &mut ConstrainedFitnessFunction<CONSTRAINTS, TChromosome, f64, TFitness, TConstraint>,
        &mut TSelection,
        &mut TPairing,
        &mut TCrossover,
        &mut TPerturbation,
        &mut TReplacement
    ) {
        let mut k_iters_feasible = VecDeque::with_capacity(k);
        move |
        iteration,
        _,
        population,
        fitness,
        _,
        _,
        _,
        _,
        _| {
            let best_candidate = population.best_candidate(better_than);
            let feasible = fitness.constraints
                .iter()
                .any(|c| c.is_feasible(&best_candidate.chromosome)
                     .expect("Can verify candidates"));

            // Change weight this iteration?
            if iteration % n == 0 {
                let all_feasible = k_iters_feasible.iter().all(|&f| f);

                for constraint_weight in fitness.constraint_weights.iter_mut() {
                    *constraint_weight *= if all_feasible {
                        1.0 / beta_1
                    } else {
                        beta_2
                    };
                }
            }

            k_iters_feasible.push_back(feasible);

            if k_iters_feasible.len() > k {
                k_iters_feasible.pop_front();
            }
        }
}

pub fn evolve_constraint_penalty_weight_tau_target
    <TChromosome: Clone,
     const CONSTRAINTS: usize,
     const DParents: usize,
     TSelection: Selection<TChromosome, f64>,
     TFitness: FitnessFunction<In = TChromosome, Out = f64>,
     TConstraint: ConstraintFunction<Chromosome = TChromosome, Out = f64>,
     TPairing: Pairing<DParents, Chromosome = TChromosome, Out = f64>,
     TCrossover: Crossover<DParents, Chromosome = TChromosome, Out = f64>,
     TReplacement: Replacement<TChromosome, f64>,
     TPerturbation: PerturbationOperator<Chromosome = TChromosome>>(
        n: usize,
        c: f64,
        tau_target: f64,
        better_than: &impl BetterThanOperator<f64>
    ) -> impl FnMut(
        usize,
        &EvolutionStats<TChromosome, f64>,
        &EvaluatedPopulation<TChromosome, f64>,

        &mut ConstrainedFitnessFunction<CONSTRAINTS, TChromosome, f64, TFitness, TConstraint>,
        &mut TSelection,
        &mut TPairing,
        &mut TCrossover,
        &mut TPerturbation,
        &mut TReplacement
    ) {
        move |
        iteration,
        _,
        population,
        fitness,
        _,
        _,
        _,
        _,
        _| {
            if iteration % n != 0 {
                return;
            }

            let count_feasible = population.population
                .iter()
                .filter(|individual| {
                    fitness.constraints
                        .iter()
                        .all(|f| f
                             .is_feasible(&individual.chromosome)
                             .expect("Can verify candidates"))
                })
                .count();
            let tau = count_feasible as f64 / population.population.len() as f64;

            for constraint_weight in fitness.constraint_weights.iter_mut() {
                if tau > tau_target {
                    *constraint_weight *= 1.0 / c;
                } else {
                    *constraint_weight *= c;
                }
            }
        }
}

            } else {
                fitness.constraints_weight *= c;
            }
        }
}