use eoa_lib::{
    crossover::Crossover,
    multi_objective_evolution::NSGAEvaluation,
    pairing,
    population::{EvaluatedChromosome, EvaluatedPopulation, Population},
};
use rand::{Rng, RngCore};

pub struct FeasibleCrossoverWrapper<
    const OBJECTIVES: usize,
    TChromosome,
    TCrossover: Crossover<2, Chromosome = TChromosome, Out = NSGAEvaluation<f64, OBJECTIVES>>,
> {
    // If there is just one infesiable, replace it
    // with this probability.
    pub p_single_replaced: f64,
    // If there are two infesiable, replace
    // first one with this probability
    pub p_double_first_replaced: f64,
    // If there are two infesiable, also
    // replace the second with this probability.
    // This is tried only if the first one is
    // replaced.
    pub p_double_second_replaced: f64,

    pub archived_count: usize,

    pub archived_population: Vec<EvaluatedChromosome<TChromosome, NSGAEvaluation<f64, OBJECTIVES>>>,
    pub crossover: TCrossover,
}

impl<
        const OBJECTIVES: usize,
        TChromosome: Clone,
        TCrossover: Crossover<2, Chromosome = TChromosome, Out = NSGAEvaluation<f64, OBJECTIVES>>,
    > FeasibleCrossoverWrapper<OBJECTIVES, TChromosome, TCrossover>
{
    pub fn update_archive(
        &mut self,
        population: &EvaluatedPopulation<TChromosome, NSGAEvaluation<f64, OBJECTIVES>>,
    ) {
        // Find all feasible solutions in population. Those are candidates
        // to replace in archive.

        // Lotta allocations...
        let mut feasible_individuals = population
            .population
            .iter()
            .filter(|individual| {
                individual
                    .evaluation
                    .evaluations
                    .iter()
                    .skip(1)
                    .all(|&constr| constr <= 0.0)
            })
            .cloned()
            .collect::<Vec<_>>();

        // TODO: this could definitely be allocated in a smarter way for better effectivity
        self.archived_population.append(&mut feasible_individuals);

        self.archived_population.sort_unstable_by(|a, b| {
            a.evaluation.evaluations[0].total_cmp(&b.evaluation.evaluations[0])
        });

        self.archived_population.truncate(self.archived_count);
    }
}

impl<
        const OBJECTIVES: usize,
        TChromosome: Clone,
        TCrossover: Crossover<2, Chromosome = TChromosome, Out = NSGAEvaluation<f64, OBJECTIVES>>,
    > Crossover<2> for FeasibleCrossoverWrapper<OBJECTIVES, TChromosome, TCrossover>
{
    type Chromosome = TChromosome;
    type Out = NSGAEvaluation<f64, OBJECTIVES>;

    fn crossover(
        &self,
        parents: &EvaluatedPopulation<Self::Chromosome, NSGAEvaluation<f64, OBJECTIVES>>,
        pairs: impl Iterator<Item = pairing::ParentPairing<2>>,
        rng: &mut dyn RngCore,
    ) -> Population<Self::Chromosome> {
        // Lotta allocations! :(

        let parents_count = parents.population.len();
        let mut joined_population = parents.clone();
        joined_population.join(EvaluatedPopulation::from_vec(
            self.archived_population.clone(),
        ));
        let full_population = joined_population.population.len();

        let mut new_pairs = pairs.collect::<Vec<_>>();

        for pair in new_pairs.iter_mut() {
            let a = &joined_population.population[pair[0]];
            let b = &joined_population.population[pair[1]];

            let a_feasible = a
                .evaluation
                .evaluations
                .iter()
                .skip(1)
                .all(|&constr| constr <= 0.0);
            let b_feasible = b
                .evaluation
                .evaluations
                .iter()
                .skip(1)
                .all(|&constr| constr <= 0.0);

            // Only proceed with replacement if we have archived feasible solutions
            if full_population > parents_count {
                match (a_feasible, b_feasible) {
                    (false, true) => {
                        if rng.random_bool(self.p_single_replaced) {
                            pair[0] = rng.random_range(parents_count..full_population);
                        }
                    }
                    (true, false) => {
                        if rng.random_bool(self.p_single_replaced) {
                            pair[1] = rng.random_range(parents_count..full_population);
                        }
                    }
                    (false, false) => {
                        if rng.random_bool(self.p_double_first_replaced) {
                            pair[0] = rng.random_range(parents_count..full_population);

                            if rng.random_bool(self.p_double_second_replaced) {
                                pair[1] = rng.random_range(parents_count..full_population);
                            }
                        }
                    }
                    (true, true) => {
                        // Do nothing.
                    }
                }
            }
        }

        self.crossover
            .crossover(&joined_population, new_pairs.into_iter(), rng)
    }
}