use std::ops::{AddAssign, Sub};
use std::{cmp::Ordering, error::Error};
use std::fmt::Debug;

use rand::RngCore;

use crate::population::EvaluatedChromosome;
use crate::{comparison::{BetterThanOperator, MinimizingOperator}, crossover::Crossover, evolution::{evolution_algorithm, EvolutionResult}, fitness::FitnessFunction, pairing::Pairing, perturbation::PerturbationOperator, population::Population, replacement::BestReplacement, selection::TournamentSelection};

// Multi objective fitness function, returning a list of evaluations
pub struct MOFitness<'a,
                             const OBJECTIVES: usize,
                             TChromosome,
                             TOut,
                             TErr: Error> {
    objectives: [Box<dyn FitnessFunction<In = TChromosome, Out = TOut, Err = TErr> + 'a>; OBJECTIVES]
}

pub fn a_dominates_b<const OBJECTIVES: usize, TOut>(
    a: &[TOut; OBJECTIVES], b: &[TOut; OBJECTIVES],
    better_than: &(impl BetterThanOperator<TOut> + ?Sized)
) -> bool {
    a.iter().zip(b.iter()).all(|(a, b)| better_than.better_than(a, b) || better_than.equal(a, b))
}

pub fn a_dominated_by_b<const OBJECTIVES: usize, TOut>(
    a: &[TOut; OBJECTIVES],
    b: &[TOut; OBJECTIVES],
    better_than: &(impl BetterThanOperator<TOut> + ?Sized)
) -> bool {
    a_dominates_b(b, a, better_than)
}

impl<'a,
     const OBJECTIVES: usize,
     TChromosome,
     TOut,
     TErr: Error>
    MOFitness<'a, OBJECTIVES, TChromosome, TOut, TErr>
{
    pub fn new(objectives: [Box<dyn FitnessFunction<In = TChromosome, Out = TOut, Err = TErr> + 'a>; OBJECTIVES]) -> Self {
        Self { objectives }
    }
}

impl<'a,
     const OBJECTIVES: usize,
     TChromosome,
     TOut: Default + Copy,
     TErr: Error + 'static>
    FitnessFunction for MOFitness<'a, OBJECTIVES, TChromosome, TOut, TErr> {
    type In = TChromosome;
    type Out = [TOut; OBJECTIVES];
    type Err = TErr;

    fn fit(&self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        let mut evaluations = [Default::default(); OBJECTIVES];

        for (i, objective) in self.objectives.iter().enumerate() {
            evaluations[i] = objective.fit(inp)?;
        }

        Ok(evaluations)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct NSGAEvaluation<TOut: PartialEq + Clone, const OBJECTIVES: usize> {
    evaluations: [TOut; OBJECTIVES],
    nondominated_front: usize,
    crowding_distance: f64
}

impl<TOut: PartialEq + Clone + PartialOrd,
     const OBJECTIVES: usize>
    PartialOrd for NSGAEvaluation<TOut, OBJECTIVES> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        match self.nondominated_front.partial_cmp(&other.nondominated_front) {
            Some(core::cmp::Ordering::Equal) => {}
            ord => return ord,
        }

        other.crowding_distance.partial_cmp(&self.crowding_distance)
    }
}

pub struct NSGAFitness<'a, const OBJECTIVES: usize, TChromosome, TOut, TErr: Error> {
    mo_fitness: MOFitness<'a, OBJECTIVES, TChromosome, TOut, TErr>,
    better_than: &'a dyn BetterThanOperator<TOut>
}

pub fn get_nondominated_fronts<const OBJECTIVES: usize,TChromosome, TOut>(
    fitted: &[EvaluatedChromosome<TChromosome, [TOut; OBJECTIVES]>],
    better_than: &(impl BetterThanOperator<TOut> + ?Sized)
) -> (Vec<usize>, Vec<Vec<usize>>) {
    let mut remaining_indices = (0..fitted.len()).collect::<Vec<_>>();
    let mut current_front = 0;
    let mut fronts = vec![0; fitted.len()];
    let mut front_indices = vec![];
    let mut current_for_removal = vec![];

    while !remaining_indices.is_empty() {
        let mut current_front_indices = vec![];
        'outer: for (i, &current) in remaining_indices.iter().enumerate() {
            let current_eval = &fitted[current].evaluation;
            for &i in remaining_indices.iter() {
                let i_eval = &fitted[i].evaluation;

                if i == current {
                    continue;
                }

                if current_eval
                    .iter()
                    .zip(i_eval.iter())
                    .all(|(a, b)| better_than.equal(a, b)) {
                    continue;
                }

                if a_dominated_by_b(current_eval, i_eval, better_than) {
                    // At least one dominates current...
                    continue 'outer;
                }
            }

            // None dominates current...
            fronts[current] = current_front;
            current_front_indices.push(current);
            current_for_removal.push(i);
        }

        front_indices.push(current_front_indices);

        for &for_removal in current_for_removal.iter().rev() {
            remaining_indices.swap_remove(for_removal);
        }
        current_for_removal.clear();

        current_front += 1;
    }

    // current_front is counting from backwards, so reverse it
    (fronts, front_indices)
}

pub fn get_front_crowding_distances<const NORMALIZE: bool,
                              const OBJECTIVES: usize,
                              TChromosome,
                              TOut: Copy + Into<f64>>(
    fitted: &[EvaluatedChromosome<TChromosome, [TOut; OBJECTIVES]>],
    indices: &[usize],
    better_than: &(impl BetterThanOperator<TOut> + ?Sized),
    crowding_distances: &mut [f64]
) {
    let mut min_maxs_init = [(Default::default(), Default::default()); OBJECTIVES];
    for objective in 0..OBJECTIVES {
        min_maxs_init[objective].0 = fitted[indices[0]].evaluation[objective].into();
        min_maxs_init[objective].1 = fitted[indices[0]].evaluation[objective].into();
    }

    // For each objective, find minimum and maximum
    let min_maxs = indices
        .iter()
        .map(|&i| &fitted[i])
        .fold(
            min_maxs_init,
            |mut min_maxs, current| {
                for objective in 0..OBJECTIVES {
                    let mut min_max = min_maxs[objective];
                    let evaluation = current.evaluation[objective].into();

                    if evaluation > min_max.1 {
                        min_max.1 = evaluation;
                    }
                    if evaluation < min_max.0 {
                        min_max.0 = evaluation;
                    }

                    min_maxs[objective] = min_max;
                }

                min_maxs
            });

    let crowding_orderings = (0..OBJECTIVES)
        .map(|objective| {
            let mut indices = indices.iter().copied().collect::<Vec<_>>();
            indices.sort_unstable_by(|&a, &b| better_than.ordering(
                &fitted[a].evaluation[objective],
                &fitted[b].evaluation[objective]
            ));
            indices
        })
        .collect::<Vec<Vec<usize>>>();

    for objective in 0..OBJECTIVES {
        let crowding_orderings = &crowding_orderings[objective];
        let mut greaters = vec![f64::INFINITY; fitted.len()];
        let mut lessers = vec![-f64::INFINITY; fitted.len()];

        let mut greaters_index = 0;
        for i in 1..indices.len() {
            let current = fitted[crowding_orderings[i]].evaluation[objective].into();
            let previous = fitted[crowding_orderings[i - 1]].evaluation[objective].into();

            if current > previous {
                lessers[i] = previous;

                while greaters_index < i {
                    greaters[greaters_index] = current;
                    greaters_index += 1;
                }
            } else {
                lessers[i] = lessers[i - 1];
            }
        }

        for (i, &current) in crowding_orderings.iter().enumerate() {
            crowding_distances[current] +=
                (greaters[i] - lessers[i]).abs() /
                if NORMALIZE { (min_maxs[objective].1 - min_maxs[objective].0).abs() } else { 1.0 };
        }
    }
}

pub fn get_crowding_distances<const NORMALIZE: bool,
                              const OBJECTIVES: usize,
                              TChromosome,
                              TOut: Copy + Into<f64>>(
    fitted: &[EvaluatedChromosome<TChromosome, [TOut; OBJECTIVES]>],
    front_indices: &[Vec<usize>],
    better_than: &(impl BetterThanOperator<TOut> + ?Sized)
) -> Vec<f64> {
    let mut crowding_distances = vec![0.0; fitted.len()];

    for indices in front_indices {
        get_front_crowding_distances::<NORMALIZE, OBJECTIVES, TChromosome, TOut>(
            fitted,
            &indices,
            better_than,
            &mut crowding_distances
        );
    }

    crowding_distances
}

impl<'a,
     const OBJECTIVES: usize,
     TChromosome,
     TOut: Copy + Into<f64>,
     TErr: Error>
    NSGAFitness<'a, OBJECTIVES, TChromosome, TOut, TErr> {
        pub fn new(
            objectives: [Box<dyn FitnessFunction<In = TChromosome, Out = TOut, Err = TErr> + 'a>; OBJECTIVES],
            better_than: &'a impl BetterThanOperator<TOut>
        ) -> Self {
            Self {
                mo_fitness: MOFitness::<'a, OBJECTIVES, _, _, _>::new(objectives),
                better_than
            }
        }
}

impl<'a,
     const OBJECTIVES: usize,
     TChromosome: Clone,
     TOut: PartialEq + Default + Copy + PartialOrd + Into<f64>,
     TErr: Error + 'static>
    FitnessFunction for NSGAFitness<'a, OBJECTIVES, TChromosome, TOut, TErr> {
    type In = TChromosome;
    type Out = NSGAEvaluation<TOut, OBJECTIVES>;
    type Err = TErr;

    fn fit(&self, _: &Self::In) -> Result<Self::Out, Self::Err> {
        panic!("NSGA fitness does not implement single fit. To get all the non dominated fronts informations, fit_population needs to be called.");
    }

    fn fit_population(&self, inp: Vec<Self::In>) -> Result<Vec<crate::population::EvaluatedChromosome<Self::In, Self::Out>>, Self::Err> {
        let fitted = self.mo_fitness.fit_population(inp)?;

        let (fronts, front_indices) = get_nondominated_fronts(&fitted, self.better_than);
        let cd = get_crowding_distances
            ::<true, OBJECTIVES, TChromosome, TOut>(&fitted, &front_indices, self.better_than);

        Ok(fitted
           .into_iter()
           .zip(fronts.into_iter())
           .zip(cd.into_iter())
           .map(|((individual, nondominated_front), crowding_distance)| {
            EvaluatedChromosome {
                chromosome: individual.chromosome,
                evaluation: NSGAEvaluation {
                    evaluations: individual.evaluation,
                    nondominated_front,
                    crowding_distance
                }
            }
           })
           .collect())
    }
}

pub fn nsga_2<const OBJECTIVES: usize,
              TChromosome: Clone,
              TResult: Clone + Debug + PartialEq + Default + Copy + PartialOrd + Into<f64>,
              TErr: Error + 'static,
              const DParents: usize,
              TFitness: FitnessFunction<In = TChromosome, Out = TResult>,
              TPairing: Pairing<DParents, Chromosome = TChromosome, Out = NSGAEvaluation<TResult, OBJECTIVES>>,
              TCrossover: Crossover<DParents, Chromosome = TChromosome, Out = NSGAEvaluation<TResult, OBJECTIVES>>,
              TPerturbation: PerturbationOperator<Chromosome = TChromosome>> (
    initial_population: Population<TChromosome>,
    parents_count: usize,
    // TODO: possibility for objectives with different outputs?
    objectives: [Box<dyn FitnessFunction<In = TChromosome, Out = TResult, Err = TErr> + '_>; OBJECTIVES],
    pairing: &mut TPairing,
    crossover: &mut TCrossover,
    perturbation: &mut TPerturbation,
    // TODO: possibility for better_than different for different fitness functions
    better_than: &impl BetterThanOperator<TResult>,
    // TODO: termination condition
    iterations: usize,
    rng: &mut dyn RngCore
) -> Result<EvolutionResult<TChromosome, [TResult; OBJECTIVES]>, Box<dyn Error>> {
    let result = evolution_algorithm(
        initial_population,
        parents_count,
        &mut NSGAFitness::new(objectives, better_than),
        &mut TournamentSelection::new(5, 1.0),
        pairing,
        crossover,
        perturbation,
        &mut BestReplacement::new(),
        &MinimizingOperator::new(),
        iterations,
        rng,
        |_, _, _, _, _, _, _, _, _| {
        }
    )?;

    Ok(result.map(|res| {
        res.evaluations
    }))
}

#[cfg(test)]
pub mod test {
    use std::str::FromStr;

    use approx::assert_abs_diff_eq;

    use crate::{comparison::MinimizingOperator, multi_objective_evolution::{get_crowding_distances, get_nondominated_fronts}, population::EvaluatedChromosome, test_infra::load_test_file};

    #[derive(Debug)]
    pub struct EvaluationResult {
        front: usize,
        crowding_distance: f64
    }

    impl FromStr for EvaluationResult {
        type Err = <f64 as FromStr>::Err;

        fn from_str(s: &str) -> Result<Self, Self::Err> {
            let mut splitted = s.split(' ');

            let front: f64 = splitted.next().unwrap().parse()?;
            let crowding_distance: f64 = splitted.next().unwrap().parse()?;

            Ok(EvaluationResult {
                front: front as usize,
                crowding_distance
            })
        }
    }

    pub fn perform_test(file: &str) {
        let data = load_test_file::<f64, EvaluationResult>(file);

        let evaluations = data
            .iter()
            .map(|vec| [vec.inp[0], vec.inp[1]])
            .collect::<Vec<_>>();

        println!("{:?}", evaluations);

        let chromosomes = evaluations
            .into_iter()
            .enumerate()
            .map(|(i, evaluation)| EvaluatedChromosome {
            chromosome: i,
            evaluation
            })
            .collect::<Vec<_>>();

        let (fronts, front_indices) = get_nondominated_fronts(&chromosomes, &MinimizingOperator::new());
        let crowding_distances = get_crowding_distances
            ::<false, 2, usize, f64>(&chromosomes, &front_indices, &MinimizingOperator::new());

        assert_eq!(fronts.len(), chromosomes.len());
        assert_eq!(crowding_distances.len(), chromosomes.len());

        println!("{:?}", fronts);
        println!("{:?}", crowding_distances);

        for (i, test) in data.iter().enumerate() {
            println!("Checking index {}", i);
            let actual_front = fronts[i];
            let actual_crowding_distance = crowding_distances[i];

            assert_eq!(actual_front + 1, test.out.front);
            if actual_crowding_distance.is_finite() {
                assert_abs_diff_eq!(actual_crowding_distance, test.out.crowding_distance, epsilon = 0.0001);
            } else {
                assert_eq!(actual_crowding_distance, test.out.crowding_distance);
            }
        }

    }

    #[test]
    pub fn test_01() {
        perform_test("tests/multi_objective_1.txt");
    }

    #[test]
    pub fn test_02() {
        perform_test("tests/multi_objective_2.txt");
    }

    #[test]
    pub fn test_03() {
        perform_test("tests/multi_objective_3.txt");
    }

    #[test]
    pub fn test_04() {
        perform_test("tests/multi_objective_4.txt");
    }

    #[test]
    pub fn test_5() {
        perform_test("tests/multi_objective_5.txt");
    }

    #[test]
    pub fn test_6() {
        perform_test("tests/multi_objective_6.txt");
    }

    #[test]
    pub fn test_7() {
        perform_test("tests/multi_objective_7.txt");
    }

    #[test]
    pub fn test_8() {
        perform_test("tests/multi_objective_8.txt");
    }

    #[test]
    pub fn test_9() {
        perform_test("tests/multi_objective_9.txt");
    }
}