use std::error::Error;
use rand::RngCore;

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

#[derive(Clone, Debug)]
pub struct EvolutionCandidate<TInput, TResult> {
    pub evaluated_chromosome: EvaluatedChromosome<TInput, TResult>,
    pub evaluation: usize,
    pub iteration: usize
}

#[derive(Clone, Debug)]
pub struct EvolutionCandidatePopulation<TInput, TResult> {
    pub current_population: EvaluatedPopulation<TInput, TResult>,
    pub iteration: usize
}

#[derive(Clone, Debug)]
pub struct EvolutionStats<TInput, TResult> {
    pub best_candidates: Vec<EvolutionCandidate<TInput, TResult>>,
}

#[derive(Clone, Debug)]
pub struct EvolutionResult<TInput, TResult> {
    pub population: EvaluatedPopulation<TInput, TResult>,
    pub stats: EvolutionStats<TInput, TResult>,
    pub best_candidate: EvaluatedChromosome<TInput, TResult>,
    pub iterations: usize,
    pub evaluations: usize
}

pub fn evolution_algorithm
    <TChromosome: Clone,
     TResult: Clone,
     const DParents: usize,
     TSelection: Selection<TChromosome, TResult>,
     TPairing: Pairing<DParents, Chromosome = TChromosome, Out = TResult>,
     TCrossover: Crossover<DParents, Chromosome = TChromosome, Out = TResult>,
     TReplacement: Replacement<TChromosome, TResult>,
     TPerturbation: PerturbationOperator<Chromosome = TChromosome>>(
    initial_population: Population<TChromosome>,
    parents_count: usize,
    fitness: &impl FitnessFunction<In = TChromosome, Out = TResult>,
    selection: &mut TSelection,
    pairing: &mut TPairing,
    crossover: &mut TCrossover,
    perturbation: &mut TPerturbation,
    replacement: &mut TReplacement,
    better_than: &impl BetterThanOperator<TResult>,
    // TODO: termination condition
    iterations: usize,
    rng: &mut dyn RngCore,
    mut evolutionary_strategy: impl FnMut(
        usize,
        &EvolutionStats<TChromosome, TResult>,
        &EvaluatedPopulation<TChromosome, TResult>,

        &mut TSelection,
        &mut TPairing,
        &mut TCrossover,
        &mut TPerturbation,
        &mut TReplacement
    )
) -> Result<EvolutionResult<TChromosome, TResult>, Box<dyn Error>> {
    let mut current_evaluation = 0;

    let mut last_best_candidate: Option<EvolutionCandidate<TChromosome, TResult>> = None;
    let mut stats: EvolutionStats<TChromosome, TResult> = EvolutionStats {
        best_candidates: vec![]
    };

    fn get_fitness_fn<TChromosome: Clone, TResult: Clone>(
        current_evaluation: &mut usize,
        better_than: &impl BetterThanOperator<TResult>,
        fitness: &impl FitnessFunction<In = TChromosome, Out = TResult>,
        current_iteration: &usize,
        stats: &mut EvolutionStats<TChromosome, TResult>,
        last_best_candidate: &mut Option<EvolutionCandidate<TChromosome, TResult>>
    ) -> impl FnMut(&TChromosome) -> TResult {
      |chromosome| {
          let evaluation = fitness.fit(chromosome).unwrap();

          if last_best_candidate.is_none() ||
              better_than.better_than(
                  &evaluation,
                  &last_best_candidate.as_ref().unwrap().evaluated_chromosome.evaluation
              ) {
                  let previous_best = std::mem::replace(
                      last_best_candidate,
                      Some(EvolutionCandidate {
                      evaluated_chromosome: EvaluatedChromosome {
                          chromosome: chromosome.clone(),
                          evaluation: evaluation.clone(),
                      },
                      evaluation: *current_evaluation,
                      iteration: *current_iteration
                  }));

                  if let Some(previous_best) = previous_best {
                      stats.best_candidates.push(previous_best);
                  }
              }
          *current_evaluation += 1;

          evaluation
      }
    }

    let mut current_population = initial_population.evaluate_mut(
        &mut get_fitness_fn(
            &mut current_evaluation,
            better_than,
            fitness,
            &0,
            &mut stats,
            &mut last_best_candidate
        )
    );

    for iteration in 1..=iterations {
        // Selection
        let parents = selection.select(parents_count, &current_population, better_than, rng).collect::<Vec<_>>();
        let parent_pairings = pairing.pair(&current_population, parents.into_iter());

        // Crossover
        let mut offsprings = crossover.crossover(&current_population, parent_pairings, rng);

        // Mutation
        for offspring in offsprings.iter_mut() {
            perturbation.perturb(offspring, rng);
        }

        let evaluated_offsprings = offsprings.evaluate_mut(
            &mut get_fitness_fn(
                &mut current_evaluation,
                better_than,
                fitness,
                &iteration,
                &mut stats,
                &mut last_best_candidate
            )
        );

        // Replace
        current_population = replacement.replace(current_population, evaluated_offsprings, better_than, rng);

        evolutionary_strategy(
            iteration,
            &stats,
            &current_population,
            selection,
            pairing,
            crossover,
            perturbation,
            replacement
        );
    }

    let best_candidate = last_best_candidate.as_ref().unwrap().evaluated_chromosome.clone();
    stats.best_candidates.push(last_best_candidate.unwrap());

    Ok(EvolutionResult {
        population: current_population,
        best_candidate,
        stats,
        iterations,
        evaluations: current_evaluation
    })
}

#[cfg(test)]
pub mod tests {
    use nalgebra::Const;

    use crate::{binary_string::BinaryString, comparison::MinimizingOperator, crossover::BinaryOnePointCrossover, fitness::one_max::OneMax, initializer::{Initializer, RandomInitializer}, pairing::AdjacentPairing, perturbation::{BinaryStringBitPerturbation, BinaryStringFlipPerturbation, BinaryStringSingleBitPerturbation, CombinedPerturbation, MutationPerturbation}, population::Population, replacement::{BestReplacement, TournamentReplacement}, selection::TournamentSelection};

    use super::evolution_algorithm;

    #[test]
    pub fn test_evolution_one_max() {
        const D: usize = 512;
        let optimum = BinaryString::<Const<D>>::new(vec![0; D]);
        let one_max = OneMax::<Const<D>>::new();

        let initializer = RandomInitializer::<Const<D>, BinaryString::<Const<D>>>::new_binary();
        let population_size = 10;

        let mut rng_init = rand::rng();
        let population = Population::from_vec(
            initializer.initialize(Const::<D>, population_size, &mut rng_init)
        );

        let mut rng = rand::rng();
        let result = evolution_algorithm(
            population,
            50,
            &one_max,
            &mut TournamentSelection::new(5, 0.8),
            &mut AdjacentPairing::new(),
            &mut BinaryOnePointCrossover::new(),
            &mut CombinedPerturbation::new(
                vec![
                    Box::new(MutationPerturbation::new(
                        Box::new(BinaryStringSingleBitPerturbation::new()),
                        0.3)),
                    Box::new(MutationPerturbation::new(
                        Box::new(BinaryStringFlipPerturbation::new()),
                        0.3))
                ]
            ),
            &mut BestReplacement::new(),
            &MinimizingOperator,
            1000,
            &mut rng,
            |_, _, _, _, _, _, _, _| ()
        ).unwrap();

        assert_eq!(
            result.best_candidate.evaluation,
            0
        );

        assert_eq!(
            result.best_candidate.chromosome,
            optimum
        );

    }
}