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>(
    initial_population: Population<TChromosome>,
    parents_count: usize,
    fitness: &impl FitnessFunction<In = TChromosome, Out = TResult>,
    selection: &impl Selection<TChromosome, TResult>,
    pairing: &mut impl Pairing<Chromosome = TChromosome, Out = TResult>,
    crossover: &impl Crossover<Chromosome = TChromosome, Out = TResult>,
    perturbation: &impl PerturbationOperator<Chromosome = TChromosome>,
    replacement: &impl Replacement<TChromosome, TResult>,
    better_than: &impl BetterThanOperator<TResult>,
    // TODO: termination condition
    iterations: usize,
    rng: &mut dyn RngCore,
) -> 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);
    }

    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,
            &TournamentSelection::new(5, 0.8),
            &mut AdjacentPairing::new(),
            &BinaryOnePointCrossover::new(),
            &CombinedPerturbation::new(
                vec![
                    Box::new(MutationPerturbation::new(
                        Box::new(BinaryStringSingleBitPerturbation::new()),
                        0.3)),
                    Box::new(MutationPerturbation::new(
                        Box::new(BinaryStringFlipPerturbation::new()),
                        0.3))
                ]
            ),
            &TournamentReplacement::new(5, 0.8),
            &MinimizingOperator,
            1000,
            &mut rng
        ).unwrap();

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

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

    }
}