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

use rand::{Rng, RngCore};
use thiserror::Error;

use crate::{comparison::{BetterThanOperator, MinimizingOperator}, crossover::Crossover, evolution::{evolution_algorithm_best_candidate, EvolutionCandidate, EvolutionResult, EvolutionStats}, fitness::FitnessFunction, pairing::Pairing, perturbation::PerturbationOperator, population::{EvaluatedChromosome, EvaluatedPopulation, Population}, 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;
                }
            }
        }
}

#[derive(PartialEq, Clone, Debug)]
pub struct ConstrainedEvaluation<const CONSTRAINTS: usize, TOut> {
    fitness: TOut,
    constraints: [TOut; CONSTRAINTS],
    weighted_sum: TOut,
    is_feasible: bool,
}

impl<const CONSTRAINTS: usize, TOut: PartialOrd> PartialOrd for ConstrainedEvaluation<CONSTRAINTS, TOut> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.weighted_sum.partial_cmp(&other.weighted_sum)
    }
}

pub struct ConstrainedEvalFitness<'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: [TOut; CONSTRAINTS]
}

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

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

        for (i, (constraint, weight)) in self.constraints.iter().zip(self.constraint_weights.iter()).enumerate() {
            let constraint_eval = match constraint.evaluate(inp) {
                Ok(constraint) => constraint,
                Err(err) =>
                    return Err(ConstrainedFitnessErr::ConstraintErr(err))
            };
            constraints[i] = constraint_eval;

            weighted_sum += match constraint.is_feasible(inp) {
                Ok(true) => weight.clone() * constraint_eval,
                Ok(false) => {
                    is_feasible = false;
                    Default::default()
                },
                Err(err) => return Err(ConstrainedFitnessErr::ConstraintErr(err))
            };
        }

        Ok(ConstrainedEvaluation {
            fitness: fit,
            constraints,
            weighted_sum,
            is_feasible
        })
    }
}

fn stochastic_ranking_sort<const CONSTRAINTS: usize, TIn, TOut: PartialOrd + Default>(
    evaluations: &[EvaluatedChromosome<TIn, ConstrainedEvaluation<CONSTRAINTS, TOut>>],
    N: usize,
    p: f64,
    better_than: &(impl BetterThanOperator<TOut> + ?Sized),
    rng: &mut dyn RngCore
) -> Vec<usize> {
    let mut indices = (0..evaluations.len()).collect::<Vec<_>>();
    for _ in 0..N {
        for j in 0..evaluations.len()-1 {
            let u = rng.random_range(0.0..=1.0);

            let current_evaluation = &evaluations[indices[j]].evaluation;
            let next_evaluation = &evaluations[indices[j + 1]].evaluation;

            if (current_evaluation.weighted_sum == Default::default() && next_evaluation.weighted_sum == Default::default()) || u < p {
                if better_than.better_than(&next_evaluation.fitness, &current_evaluation.fitness) {
                    indices.swap(j, j + 1);
                }
            } else {
                if current_evaluation.weighted_sum > next_evaluation.weighted_sum {
                    indices.swap(j, j + 1);
                }
            }
        }
    }

    indices
}

pub struct StochasticRankingSelection<'a, TSelection, TBetterThan> {
    N: usize,
    p: f64,
    selection: &'a mut TSelection,
    better_than: &'a TBetterThan
}

const MINIMIZING_OPERATOR: MinimizingOperator = MinimizingOperator;

impl<'a,
     const CONSTRAINTS: usize,
     TChromosome,
     TResult: PartialOrd + Default,
     TSelection: Selection<(), usize>,
     TBetterThan: BetterThanOperator<TResult>>
    Selection<TChromosome, ConstrainedEvaluation<CONSTRAINTS, TResult>> for StochasticRankingSelection<'a, TSelection, TBetterThan> {
    fn select(&self,
              count: usize,
              evaluations: &EvaluatedPopulation<TChromosome, ConstrainedEvaluation<CONSTRAINTS, TResult>>,
              _: &dyn BetterThanOperator<ConstrainedEvaluation<CONSTRAINTS, TResult>>,
              rng: &mut dyn RngCore
    ) -> impl Iterator<Item = usize> {
        let sorted_indices = stochastic_ranking_sort(
            evaluations.population.as_slice(),
            self.N, self.p, self.better_than, rng
        );
        let mut rankings = vec![EvaluatedChromosome {
            chromosome: (),
            evaluation: 0
        }; evaluations.population.len()];

        for (ranking, index) in sorted_indices.into_iter().enumerate() {
            rankings[index] = EvaluatedChromosome {
                chromosome: (),
                evaluation: ranking
            };
        }

        // Replace with this better than
        self.selection.select(
            count,
            &EvaluatedPopulation::from_vec(rankings),
            &MINIMIZING_OPERATOR,
            rng)
            .collect::<Vec<_>>()
            .into_iter()
    }
}

pub fn stochastic_ranking_evolution_algorithm
    <TChromosome: Clone,
     TResult: PartialOrd + std::ops::Mul<Output = TResult> + std::ops::AddAssign + Copy + Default,
     const DParents: usize,
     const CONSTRAINTS: usize,
     TFitness: FitnessFunction<In = TChromosome, Out = TResult>,
     TSelection: Selection<(), usize>,
     TPairing: Pairing<DParents, Chromosome = TChromosome, Out = ConstrainedEvaluation<CONSTRAINTS, TResult>>,
     TCrossover: Crossover<DParents, Chromosome = TChromosome, Out = ConstrainedEvaluation<CONSTRAINTS, TResult>>,
     TReplacement: Replacement<TChromosome, ConstrainedEvaluation<CONSTRAINTS, TResult>>,
     TConstraint: ConstraintFunction<Chromosome = TChromosome, Out = TResult>,
     TPerturbation: PerturbationOperator<Chromosome = TChromosome>>(
        initial_population: Population<TChromosome>,
        parents_count: usize,
        N: usize,
        p: f64,
        fitness: &mut TFitness,
        constraints: [&TConstraint; CONSTRAINTS],
        constraint_weights: [TResult; CONSTRAINTS],
        pairing: &mut TPairing,
        selection: &mut TSelection,
        crossover: &mut TCrossover,
        perturbation: &mut TPerturbation,
        replacement: &mut TReplacement,
        better_than: &(impl BetterThanOperator<ConstrainedEvaluation<CONSTRAINTS, TResult>> + BetterThanOperator<TResult>),
        // TODO: termination condition
        iterations: usize,
        rng: &mut dyn RngCore,
        // mut evolutionary_strategy: impl FnMut(
        //     usize,
        //     &EvolutionStats<TChromosome, TResult>,
        //     &EvaluatedPopulation<TChromosome, TResult>,

        //     &mut TPairing,
        //     &mut TCrossover,
        //     &mut TPerturbation,
        //     &mut TReplacement,
        //     &mut ConstrainedEvalFitness<CONSTRAINTS, TChromosome, TResult, TFitness, TConstraint>,
        // ),
    ) -> Result<(EvolutionResult<TChromosome, TResult>, Vec<f64>), Box<dyn Error>>
{
    let mut constrained_fitness = ConstrainedEvalFitness {
        fitness,
        constraints,
        constraint_weights,
    };

    let mut stochastic_ranking_selection =
        StochasticRankingSelection {
            N,
            p,
            selection,
            better_than
        };

    let mut feasible_fractions = Vec::with_capacity(iterations);

    evolution_algorithm_best_candidate(
        initial_population,
        parents_count,
        &mut constrained_fitness,
        &mut stochastic_ranking_selection,
        pairing,
        crossover,
        perturbation,
        replacement,
        better_than,
        iterations,
        rng,
        |iteration, stats, population, fitness, selection, pairing, crossover, perturbation, replacement| {
            let feasible_fraction = population.population
                .iter()
                .map(|x| x.evaluation.is_feasible)
                .count() as f64 / population.population.len() as f64;

            feasible_fractions.push(feasible_fraction);

            // evolutionary_strategy(
            //     iteration,
            //     stats,
            //     population,
            //     fitness
            // )
        },
        |_, evaluation, best_candidate| {
            // Do not save enfeasible solutions!
            if !evaluation.is_feasible {
                return false;
            }

            if best_candidate.is_none() {
                return true;
            }

            better_than.better_than(
                evaluation,
                &best_candidate.as_ref().unwrap().evaluated_chromosome.evaluation
            )
        }).map(|res|
               (
                   res.map(|evaluation| evaluation.fitness),
                   feasible_fractions
               ))
}