use rand::{seq::IteratorRandom, RngCore};
use std::fmt::Debug;

use crate::{comparison::BetterThanOperator, fitness::FitnessFunction, population::{EvaluatedPopulation, EvaluatedChromosome}, selection::{Selection, TournamentSelection}};

fn extract_by_indices<T>(mut x: Vec<T>, mut idxs: Vec<usize>) -> Vec<T> {
    idxs.sort_unstable_by(|a, b| b.cmp(a));

    let mut result = Vec::with_capacity(idxs.len());
    for idx in idxs {
        if idx < x.len() {
            result.push(x.swap_remove(idx));
        }
    }

    result.reverse();
    result
}

pub trait Replacement<TChromosome, TResult> {
    fn replace(
        &self,
        parents_evaluations: EvaluatedPopulation<TChromosome, TResult>,
        offsprings_evaluations: EvaluatedPopulation<TChromosome, TResult>,
        better_than: &dyn BetterThanOperator<TResult>,
        rng: &mut dyn RngCore
    ) -> EvaluatedPopulation<TChromosome, TResult>;
}

pub struct BestReplacement;
impl BestReplacement {
    pub fn new() -> Self {
        Self
    }
}

impl<TChromosome, TResult: Copy + Debug> Replacement<TChromosome, TResult> for BestReplacement {
    fn replace(
        &self,
        parents_evaluations: EvaluatedPopulation<TChromosome, TResult>,
        offsprings_evaluations: EvaluatedPopulation<TChromosome, TResult>,
        better_than: &dyn BetterThanOperator<TResult>,
        _rng: &mut dyn RngCore
    ) -> EvaluatedPopulation<TChromosome, TResult> {
        let count = parents_evaluations.population.len();
        let mut population = parents_evaluations;
        population.join(offsprings_evaluations);

        let mut idxs = (0..population.population.len())
            .collect::<Vec<_>>();
        idxs.sort_unstable_by(|&i, &j| better_than.ordering(
            &population.population[i].evaluation,
            &population.population[j].evaluation)
        );

        idxs.truncate(count);

        EvaluatedPopulation::from_vec(
            extract_by_indices(population.deconstruct(), idxs)
        )
    }
}

pub struct GenerationalReplacement;
impl<TInput, TResult> Replacement<TInput, TResult> for GenerationalReplacement {
    fn replace(
        &self,
        parents: EvaluatedPopulation<TInput, TResult>,
        mut offsprings: EvaluatedPopulation<TInput, TResult>,
        _: &dyn BetterThanOperator<TResult>,
        _rng: &mut dyn RngCore
    ) -> EvaluatedPopulation<TInput, TResult> {
        let count = parents.population.len();
        if count == offsprings.population.len() {
            return offsprings;
        }

        offsprings.join(parents);
        offsprings.population.truncate(count);
        // let population = offsprings.deconstruct();
        // population.truncate(count);

        // EvaluatedPopulation::from_vec(population)

        offsprings
    }
}

pub struct RandomReplacement;

impl RandomReplacement {
    pub fn new() -> Self {
        Self
    }
}

impl<TInput, TResult> Replacement<TInput, TResult> for RandomReplacement {
    fn replace(
        &self,
        parents: EvaluatedPopulation<TInput, TResult>,
        offsprings: EvaluatedPopulation<TInput, TResult>,
        _: &dyn BetterThanOperator<TResult>,
        rng: &mut dyn RngCore
    ) -> EvaluatedPopulation<TInput, TResult> {
        let count = parents.population.len();

        EvaluatedPopulation::from_vec(
            parents.deconstruct()
                .into_iter()
                .chain(offsprings.deconstruct().into_iter())
                .choose_multiple(rng, count))
    }
}

pub struct TournamentReplacement {
    selection: TournamentSelection,
    evaluation_pool: Vec<f64>
}

impl TournamentReplacement {
    pub fn new(k: usize, p: f64) -> Self {
        TournamentReplacement {
            evaluation_pool: vec![],
            selection: TournamentSelection::new(
                k,
                p,
            )
        }
    }
}

impl<TInput, TResult: Copy + Debug> Replacement<TInput, TResult> for TournamentReplacement {
    fn replace(
        &self,
        parents: EvaluatedPopulation<TInput, TResult>,
        offsprings: EvaluatedPopulation<TInput, TResult>,
        better_than: &dyn BetterThanOperator<TResult>,
        rng: &mut dyn RngCore
    ) -> EvaluatedPopulation<TInput, TResult> {
        let count = parents.population.len();
        let mut population = parents;
        population.join(offsprings);

        // TODO: use a pool instead of allocating vector every run of this function
        let selected = self.selection.select(count, &population, better_than, rng)
            .collect::<Vec<_>>();

        let population = population.deconstruct();
        let population = extract_by_indices(population, selected);

        EvaluatedPopulation::from_vec(population)
    }
}