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: Option<EvaluatedChromosome<TInput, TResult>>,
    pub iterations: usize,
    pub evaluations: usize
}

impl<TInput, TResult> EvolutionResult<TInput, TResult> {
    pub fn map<TNewResult>(self, map: impl Fn(TResult) -> TNewResult) -> EvolutionResult<TInput, TNewResult> {
        EvolutionResult {
            population: EvaluatedPopulation::from_vec(
                self.population.deconstruct()
                    .into_iter()
                    .map(|chromosome| EvaluatedChromosome {
                        chromosome: chromosome.chromosome,
                        evaluation: map(chromosome.evaluation)
                    })
                    .collect()
            ),
            stats: EvolutionStats {
                best_candidates: self.stats.best_candidates
                    .into_iter()
                    .map(|candidate|
                         EvolutionCandidate {
                             evaluated_chromosome: EvaluatedChromosome {
                                 chromosome: candidate.evaluated_chromosome.chromosome,
                                 evaluation: map(candidate.evaluated_chromosome.evaluation)
                             },
                             evaluation: candidate.evaluation,
                             iteration: candidate.iteration
                         })
                    .collect()
            },
            best_candidate: self.best_candidate.map(
                |best_candidate|
                EvaluatedChromosome {
                    chromosome: best_candidate.chromosome,
                    evaluation: map(best_candidate.evaluation)
                }),
            evaluations: self.evaluations,
            iterations: self.iterations,
        }
    }
}

pub fn evolution_algorithm
    <TChromosome: Clone,
     TResult: Clone,
     const DParents: usize,
     TSelection: Selection<TChromosome, TResult>,
     TFitness: FitnessFunction<In = TChromosome, Out = 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: &mut TFitness,
    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,
    evolutionary_strategy: impl FnMut(
        usize,
        &EvolutionStats<TChromosome, TResult>,
        &EvaluatedPopulation<TChromosome, TResult>,

        &mut TFitness,
        &mut TSelection,
        &mut TPairing,
        &mut TCrossover,
        &mut TPerturbation,
        &mut TReplacement
    )
) -> Result<EvolutionResult<TChromosome, TResult>, Box<dyn Error>> {
        evolution_algorithm_best_candidate(
            initial_population,
            parents_count,
            fitness,
            selection,
            pairing,
            crossover,
            perturbation,
            replacement,
            better_than,
            iterations,
            rng,
            evolutionary_strategy,
            |_, evaluation, last_best_candidate| last_best_candidate.is_none() ||
                better_than.better_than(
                    evaluation,
                    &last_best_candidate.as_ref().unwrap().evaluated_chromosome.evaluation
                ))
}

pub fn evolution_algorithm_best_candidate
    <TChromosome: Clone,
     TResult: Clone,
     const DParents: usize,
     TSelection: Selection<TChromosome, TResult>,
     TFitness: FitnessFunction<In = TChromosome, Out = 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: &mut TFitness,
    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 TFitness,
        &mut TSelection,
        &mut TPairing,
        &mut TCrossover,
        &mut TPerturbation,
        &mut TReplacement
    ),
    // For the statistics, evaluate if a candidate is better. Potential for different functrion than better_than that's used
    // for the replacement, selection etc.
    better_than_stats: impl Fn(&TChromosome, &TResult, &Option<EvolutionCandidate<TChromosome, TResult>>) -> bool,
) -> 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 apply_new_eval<TChromosome: Clone, TResult: Clone>(
        current_evaluation: &mut usize,
        current_iteration: &usize,
        stats: &mut EvolutionStats<TChromosome, TResult>,
        population: &EvaluatedPopulation<TChromosome, TResult>,
        last_best_candidate: &mut Option<EvolutionCandidate<TChromosome, TResult>>,
        better_than_stats: &impl Fn(&TChromosome, &TResult, &Option<EvolutionCandidate<TChromosome, TResult>>) -> bool,
    ) {
        for individual in population.iter() {
            let evaluation = &individual.evaluation;
            let chromosome = &individual.chromosome;

            if better_than_stats(chromosome, evaluation, last_best_candidate) {
                    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;
        }
    }

    let mut current_population =
        initial_population.evaluate(fitness)?;
    apply_new_eval(
        &mut current_evaluation,
        &0,
        &mut stats,
        &current_population,
        &mut last_best_candidate,
        &better_than_stats);

    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(fitness)?;
        apply_new_eval(
            &mut current_evaluation,
            &iteration,
            &mut stats,
            &evaluated_offsprings,
            &mut last_best_candidate,
            &better_than_stats
        );

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

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

    let best_candidate = last_best_candidate.as_ref().map(|x| x.evaluated_chromosome.clone());
    if last_best_candidate.is_some() {
        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 mut 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,
            &mut 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
        );

    }
}