use std::cmp::Ordering;

use genetisk::{Individual, Simulation, Wrapper};
use rand::distributions::{IndependentSample, Range};
use rayon::prelude::*;

#[derive(Clone, Copy, Debug)]
struct Fitness(f64);

impl PartialEq for Fitness {
    fn eq(&self, other: &Fitness) -> bool {
        (self.0 - other.0).abs() < 0.0001
    }
}

impl Eq for Fitness {}

impl PartialOrd for Fitness {
    fn partial_cmp(&self, other: &Fitness) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl Ord for Fitness {
    fn cmp(&self, other: &Fitness) -> Ordering {
        self.0.partial_cmp(&other.0).unwrap_or(Ordering::Equal)
    }
}

#[derive(Clone, Debug)]
struct Parabole {
    x: f64,
}

impl Individual for Parabole {
    type Fitness = Fitness;

    fn mate(&self, other: &Parabole) -> Parabole {
        Parabole {
            x: (self.x + other.x) / 2.0,
        }
    }

    fn mutate(&mut self) {
        let mut rng = rand::weak_rng();

        let between = Range::new(-1.0, 1.0);
        let offset = between.ind_sample(&mut rng);

        self.x += offset;
    }

    fn fitness(&self) -> Self::Fitness {
        Fitness(10.0 - ((self.x + 3.0) * (self.x + 3.0)))
    }
}

fn main() {
    let initial = (0..300)
        .map(|i| Parabole { x: i as f64 })
        .collect::<Vec<_>>();

    let population_size = initial.len();

    let mut simulation = Simulation::with_population(initial);

    simulation.calculate();

    for _ in 0..250_000 {
        simulation.evolve(|parents, population| {
            // Mate top 10 to get 5 children.
            parents
                .iter()
                .enumerate()
                .filter(|&(n, _)| n % 2 == 0)
                .map(|(_, wrapper)| wrapper)
                .take(5)
                .zip(
                    parents
                        .iter()
                        .enumerate()
                        .filter(|&(n, _)| n % 2 == 1)
                        .map(|(_, wrapper)| wrapper)
                        .take(5),
                )
                .map(|(a, b)| a.individual.mate(&b.individual))
                .map(|individual| Wrapper {
                    individual,
                    fitness: None,
                })
                .for_each(|wrapper| {
                    population.push(wrapper);
                });

            // Mutate all to get new children.
            parents
                .par_iter()
                .map(|wrapper| wrapper.individual.clone())
                .map(|mut individual| {
                    individual.mutate();

                    Wrapper {
                        individual,
                        fitness: None,
                    }
                })
                .collect_into_vec(population);

            // Add all parents again.
            population.extend(parents.iter().cloned());
        });

        simulation.population.truncate(population_size);
    }

    println!(
        "{:#?}",
        simulation
            .population
            .iter()
            .take(10)
            .map(|individual| &individual.individual)
            .collect::<Vec<_>>()
    );
    println!();
}