//! Common genetic algorithm.

extern crate rand;
extern crate rayon;
extern crate serde;

use rand::prelude::*;
use rayon::prelude::*;
use serde::{Serialize, Deserialize};

use std::marker::Sync;

use crate::{
    bitransform::Bitransform,
    boolnet::Boolnet
};

#[derive(Serialize, Deserialize)]
pub struct Parameters {
    pub batches: usize,
    pub replace_ratio: f64,
    pub par_ratio: f64,
    pub parents: usize,
    pub par_sw_prob: f64,
    pub mutat_prob: f64
}

pub struct Statistics {
    pub age: (u64, f64, u64) // min, average, max
}

trait Individual {
    fn test<T: Bitransform + ?Sized>(&mut self, bitran: &T, input_batches: & Vec<Vec<u128>>);
    fn age(&self) -> u64;
    fn inc_age(&mut self);
}

impl Individual for Boolnet {
    fn test<T: Bitransform + ?Sized>(&mut self, bitran: &T, input_batches: & Vec<Vec<u128>>) {
        let mut score_sum: f64 = 0.0;

        let mut outp = vec![0u128; bitran.outp_size()];

        let batches = input_batches.len();

        let mut rng = thread_rng();

        for i in 0..batches {
            let inp = & input_batches[i];

            self.run_batch(&mut rng, inp, &mut outp);

            score_sum += bitran.score_inp_outp(inp, &outp);
        }

        self.set_score(score_sum / (batches as f64));
    }

    fn age(&self) -> u64 {
        self.age()
    }

    fn inc_age(&mut self) {
        self.inc_age();
    }
}

fn make_input_batches<T: Bitransform + Sync + ?Sized>(bitran: &T, batches: usize) -> Vec<Vec<u128>> {
    let mut input_batches = vec![];
    input_batches.par_extend((0..batches).into_par_iter().map(|_| { // parallelized
        let mut rng = thread_rng();
        bitran.rand_inp(&mut rng)
    }));
    input_batches
}

pub trait Population {
    fn stats(&self) -> Statistics;
    fn score_stats(&self) -> (f64, f64, f64);
    fn unscore(&mut self);
    fn add_inheritants(&mut self, par_ratio: f64, parents: usize, par_sw_prob: f64, mutat_prob: f64, num: usize);
    fn evolearn<T: Bitransform + Sync + ?Sized>(&mut self, bitran: &T, params: &Parameters, test_all: bool);
}

impl Population for Vec<Boolnet> {
    /// Returns min, average, max age
    fn stats(&self) -> Statistics {
        let (mut min_age, mut avg_age, mut max_age) = (std::u64::MAX, 0.0f64, std::u64::MIN);
        for bn in self {
            let age = bn.age();
            min_age = min_age.min(age);
            avg_age += age as f64;
            max_age = max_age.max(age);
        }
        avg_age /= self.len() as f64;
        Statistics{age: (min_age, avg_age, max_age)}
    }

    /// Returns min, average, max scores
    /// Population must be sorted by score, descending
    fn score_stats(&self) -> (f64, f64, f64) {
        let size = self.len();
        if size > 0 {
            (
                self[size - 1].score(),
                self.iter().fold(0.0f64, |acc, bn| acc + bn.score()) / (size as f64),
                self[0].score()
            )
        } else {
            (std::f64::NEG_INFINITY, std::f64::NEG_INFINITY, std::f64::NEG_INFINITY)
        }
    }

    fn unscore(&mut self) {
        for i in 0..self.len() {
            self[i].set_score(std::f64::NEG_INFINITY);
        }
    }

    fn add_inheritants(&mut self, par_ratio: f64, parents: usize, par_sw_prob: f64, mutat_prob: f64, num: usize) {
        // Assertion: population is sorted by score, descending, so that only "best" par_ratio part of population gets descendants
        let par_size = ((par_ratio * (self.len() as f64)).round() as usize).min(self.len() - 1).max(1);

        let mut inheritants = vec![];
        inheritants.par_extend((0..num).into_par_iter().map(|_| { // parallelized
            let mut rng = thread_rng();
            let mut boolnet_parents = Vec::<&Boolnet>::with_capacity(parents);
            (0..parents).for_each( |_| boolnet_parents.push(& self[rng.gen_range(0, par_size)]) );
            Boolnet::new_inherit(&mut rng, boolnet_parents, par_sw_prob, mutat_prob)
        }));
        self.extend(inheritants);
    }

    fn evolearn<T: Bitransform + Sync + ?Sized>(&mut self, bitran: &T, params: &Parameters, test_all: bool) {
        let p_total = self.len();
        let p_replace = ((p_total as f64) * params.replace_ratio) as usize;
        let p_keep = p_total - p_replace;

        let mut all_scored = true;
        for i in 0..self.len() {
            if self[i].score() == std::f64::NEG_INFINITY {
                all_scored = false;
                break;
            }
        }

        if all_scored { // and sorted... see below
            // Remove the lowest scored ones
            self.truncate(p_keep);

            // Add inheritants of the highest scored ones
            self.add_inheritants(params.par_ratio, params.parents, params.par_sw_prob, params.mutat_prob, p_replace);
        }

        // Make input batches
        let input_batches = make_input_batches(bitran, params.batches);

        // Score
        self.into_par_iter().for_each(|bn| { // parallelized with Rayon iterator
            if test_all || (bn.score() == std::f64::NEG_INFINITY) {
                bn.test(bitran, &input_batches);
            }
            bn.inc_age();
        });

        // Sort
        self.sort_by(|a, b| b.score().partial_cmp(& a.score()).unwrap_or(std::cmp::Ordering::Equal));
        // Now the highest scored ones are in the beginning
    }

}
