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

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

#[derive(Clone, Debug)]
pub struct EvolutionCandidate<TInput, TResult> {
    pub evaluated_chromosome: EvaluatedChromosome<TInput, TResult>,
    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 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 evolution: 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_population = initial_population.evaluate(fitness)?;

    let mut last_best_candidate = EvolutionCandidate {
        evaluated_chromosome: current_population.best_candidate(better_than).clone(),
        iteration: 0
    };
    let mut stats: EvolutionStats<TChromosome, TResult> = EvolutionStats {
        best_candidates: vec![]
    };

    for iteration in 0..iterations {
        // Figure out best candidate and save it if better than last time
        let best_candidate = current_population.best_candidate(better_than);

        if better_than.better_than(
            &best_candidate.evaluation,
            &last_best_candidate.evaluated_chromosome.evaluation
        ) {
            stats.best_candidates.push(last_best_candidate);

            last_best_candidate = EvolutionCandidate {
                evaluated_chromosome: best_candidate.clone(),
                iteration
            }
        }

        // 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)?;

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

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

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

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

#[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}, replacement::{BestReplacement, Population, 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
        );

    }
}