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

use crate::{comparison::BetterThanOperator, replacement::EvaluatedPopulation};

pub trait Selection<TChromosome, TResult> {
    fn select(&mut self,
              count: usize,
              evaluations: &EvaluatedPopulation<TChromosome, TResult>,
              better_than: &dyn BetterThanOperator<TResult>
    ) -> impl Iterator<Item = usize>;
}

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

impl<TChromosome, TResult: Copy> Selection<TChromosome, TResult> for BestSelection {
    fn select(&mut self,
              count: usize,
              evaluations: &EvaluatedPopulation<TChromosome, TResult>,
              better_than: &dyn BetterThanOperator<TResult>
    ) -> impl Iterator<Item = usize> {
        let mut idxs = (0..evaluations.population.len())
            .collect::<Vec<_>>();
        idxs.sort_unstable_by(|&i, &j| better_than.ordering(
            &evaluations.population[i].evaluation,
            &evaluations.population[j].evaluation)
        );

        idxs.into_iter().take(count)
    }
}

pub struct TournamentSelection {
    rng: Box<dyn RngCore>,
    p: f64,
    k: usize
}

impl TournamentSelection {
    pub fn new(k: usize, p: f64) -> Self {
        assert!(0.0 <= p && p <= 1.0);
        assert!(k > 0);

        Self {
            rng: Box::new(rand::rng()),
            p,
            k
        }
    }

    fn tournament<TChromosome, TResult: Debug + Copy>(
        &mut self,
        idxs: &mut Vec<usize>,
        evaluations: &EvaluatedPopulation<TChromosome, TResult>,
        better_than: &dyn BetterThanOperator<TResult>
    ) -> usize {
        idxs.sort_unstable_by(|&i, &j| better_than.ordering(
            &evaluations.population[i].evaluation,
            &evaluations.population[j].evaluation)
        );

        let mut p_selector = self.rng.random_range(0.0..=1.0f64);
        let p = self.p;
        let k = self.k;

        let mut selected = idxs[k - 1];
        // let's say p = 0.7
        // the best has probability 0.7 of being selected
        // if the best is not selected, the second has 0.7 probability of being selected... (that's 0.7 * 0.3 without conditions)
        // and so on. The last element is selected if none of the previous were.
        for i in 0..k-1 {
            if p_selector <= p {
                selected = i;
                break;
            }

            p_selector -= p;
            // 'Expand' the rest to '100%' again
            p_selector /= 1.0 - p;
        }

        idxs[selected]
    }
}

impl<TChromosome, TResult: Copy + Debug> Selection<TChromosome, TResult> for TournamentSelection {
    fn select(
        &mut self,
        count: usize,
        evaluations: &EvaluatedPopulation<TChromosome, TResult>,
        better_than: &dyn BetterThanOperator<TResult>
    ) -> impl Iterator<Item = usize> {
        // Let's reuse a single vector for the indices
        let mut k_selected_idxs = vec![0; self.k];
        (0..count).map(move |_| {
            // Choose k. Do not care if already selected previously.
            (0..evaluations.population.len()).choose_multiple_fill(&mut self.rng, &mut k_selected_idxs);
            // Tournament between the k
            let index = self.tournament(&mut k_selected_idxs, evaluations, better_than);
            index
        })
    }
}