use nalgebra::Scalar;
use rand::{seq::IteratorRandom, Rng, RngCore};
use rand_distr::uniform::SampleUniform;
use std::{cmp::Ordering, fmt::Debug, ops::{AddAssign, Sub}};

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

pub trait Selection<TChromosome, TResult> {
    fn select(&self,
              count: usize,
              evaluations: &EvaluatedPopulation<TChromosome, TResult>,
              better_than: &dyn BetterThanOperator<TResult>,
              rng: &mut dyn RngCore
    ) -> 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(&self,
              count: usize,
              evaluations: &EvaluatedPopulation<TChromosome, TResult>,
              better_than: &dyn BetterThanOperator<TResult>,
              _rng: &mut dyn RngCore
    ) -> 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 {
    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 {
            p,
            k
        }
    }

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

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

        let mut selected = 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: Debug> Selection<TChromosome, TResult> for TournamentSelection {
    fn select(
        &self,
        count: usize,
        evaluations: &EvaluatedPopulation<TChromosome, TResult>,
        better_than: &dyn BetterThanOperator<TResult>,
        rng: &mut dyn RngCore
    ) -> 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(rng, &mut k_selected_idxs);
            // Tournament between the k
            let index = self.tournament(&mut k_selected_idxs, evaluations, better_than, rng);
            index
        })
    }
}


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

impl<TChromosome,
     TResult: Scalar + Copy + Default + PartialOrd + SampleUniform + AddAssign + Sub<Output = TResult>>
    Selection<TChromosome, TResult> for RouletteWheelSelection
{
    fn select(
        &self,
        count: usize,
        evaluations: &EvaluatedPopulation<TChromosome, TResult>,
        _: &dyn BetterThanOperator<TResult>,
        rng: &mut dyn RngCore
    ) -> impl Iterator<Item = usize> {
        let zero: TResult = Default::default();
        let max = evaluations.iter()
            .map(|i| i.evaluation)
            .max_by(|a, b|
                    a.partial_cmp(b).unwrap_or(Ordering::Less))
            .unwrap();
        let summed = evaluations.iter().scan(
            zero,
            |acc, individual| {
                let subtracted: TResult = max - individual.evaluation;
                *acc += subtracted;
                Some(*acc)
            })
            .collect::<Vec<TResult>>();
        let max = summed.last().unwrap().clone();

        (0..count).map(move |_| {
            let rand = rng.random_range(zero..=max);

            (0..summed.len())
                .filter(|&i| summed[i] > rand)
                .next()
                .unwrap_or(summed.len() - 1)
        })
    }
}