mod arch_union_popul;
mod properties;

use std::{fmt::Debug, ops::AddAssign};

use mop_blocks::mp::mphos::{Morhos, Mphos, MphosEvaluators};
use mop_blocks::ObjDirection;
use mop_solver::Solver;
use num_traits::Float;

use etc::{euclidean_distance, verify_pareto_dominance, ParetoDominance};
use genetic_algorithm::{
    operators::{
        crossover::Crossover, mating_selection::MatingSelection, mutation::Mutation,
        survivor_selection::SurvivorSelection,
    },
    spea2::{arch_union_popul::ArchUnionPopul, properties::Properties},
    GeneticAlgorithmParams,
};

#[derive(Clone, Debug)]
pub struct Spea2<'a, C, CR, MA, MU, N, S, V> {
    arch_mp: Morhos<N, V>,
    arch_u_popul: ArchUnionPopul<N, V>,
    archive_size: usize,
    gap: GeneticAlgorithmParams<CR, MA, MU, S>,
    k: usize,
    popul_mp: Mphos<'a, C, N, V>,
}

impl<'a, C, CR, MA, MU, N, S, V> Spea2<'a, C, CR, MA, MU, N, S, V>
where
    N: AddAssign + Copy + Debug + Float + Send + Sync,
    V: Copy + Debug + Send + Sync,
{
    pub fn new(
        archive_size: f64,
        gap: GeneticAlgorithmParams<CR, MA, MU, S>,
        mp: Mphos<'a, C, N, V>,
    ) -> Self {
        let population_size = mp.results().len();
        let archive_size = (archive_size * population_size as f64) as usize;
        let arch_u_popul_len = archive_size + population_size;
        let arch_mp = Morhos::with_capacity(archive_size, mp.definitions());
        let arch_u_popul = ArchUnionPopul {
            mp: Morhos::with_capacity(arch_u_popul_len, mp.definitions()),
            props: vec![
                Properties {
                    density: N::zero(),
                    distance: N::zero(),
                    raw_fitness: N::zero(),
                    strength: N::zero(),
                };
                arch_u_popul_len
            ],
        };
        Spea2 {
            arch_mp,
            arch_u_popul,
            archive_size,
            gap,
            k: N::from(arch_u_popul_len)
                .unwrap()
                .sqrt()
                .to_usize()
                .unwrap(),
            popul_mp: mp,
        }
    }

    fn assign_fitness(&mut self) {
        self.set_strength();
        self.set_raw_fitness();
        self.set_density();
    }

    // Fill archive
    fn copy_to_archive(&mut self) {
        let pop_len = self.arch_u_popul.mp.len();
        for first_ind_idx in 0..pop_len {
            let first_props = &self.arch_u_popul.props[first_ind_idx];
            let first_sol_violations: usize = self
                .arch_u_popul
                .mp
                .get(first_ind_idx)
                .hard_cstrs()
                .iter()
                .map(|x| x.weight())
                .sum();
            for second_ind_idx in first_ind_idx..pop_len {
                let second_props = &self.arch_u_popul.props[second_ind_idx];
                let second_sol_violations: usize = self
                    .arch_u_popul
                    .mp
                    .get(second_ind_idx)
                    .hard_cstrs()
                    .iter()
                    .map(|x| x.weight())
                    .sum();
                if (first_sol_violations == 0 && second_sol_violations > 0)
                    || (first_sol_violations > 0
                        && second_sol_violations > 0
                        && first_sol_violations < second_sol_violations)
                    || (first_sol_violations == 0 && second_sol_violations == 0
                        && first_props.density + first_props.raw_fitness
                            < second_props.density + second_props.raw_fitness)
                    || first_props.density + first_props.raw_fitness < N::one()
                {
                    self.arch_mp
                        .append_result(self.arch_u_popul.mp.get(first_ind_idx));
                }
            }
        }
    }

    fn environment_selection(&mut self) {
        self.copy_to_archive();
        if self.arch_mp.len() == self.archive_size {

        } else if self.arch_mp.len() < self.archive_size {
            // É preciso sortear a população
            let diff = self.archive_size - self.arch_mp.len();
            for idx in 0..diff {
                self.arch_mp.append_result(self.arch_u_popul.mp.get(idx));
            }
        } else {
            // TODO
            self.arch_mp.truncate(self.archive_size);
        }
    }

    fn set_density(&mut self) {
        let pop_len = self.arch_u_popul.mp.len();
        let mut tmp = Vec::with_capacity(pop_len);
        for first_ind_idx in 0..pop_len {
            for second_ind_idx in 0..pop_len {
                tmp.push(euclidean_distance(
                    self.arch_u_popul.mp.get(first_ind_idx).objs(),
                    self.arch_u_popul.mp.get(second_ind_idx).objs(),
                ));
            }
            tmp.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
            self.arch_u_popul.props[first_ind_idx].distance =
                N::one() / (tmp[self.k] + N::from(2).unwrap());
            tmp.clear();
        }
    }

    fn set_strength(&mut self) {
        let objs = self.popul_mp.definitions().objs();
        let mp = &self.arch_u_popul.mp;
        let props = &mut self.arch_u_popul.props;
        for (first_ind_idx, first_ind) in mp.iter().enumerate() {
            for (second_ind_idx, second_ind) in mp.iter().enumerate().skip(first_ind_idx) {
                match verify_pareto_dominance(first_ind.objs(), second_ind.objs(), |idx, a, b| {
                    match objs[idx].obj_direction() {
                        ObjDirection::Max => b.partial_cmp(a).unwrap(),
                        ObjDirection::Min => a.partial_cmp(b).unwrap(),
                    }
                }) {
                    ParetoDominance::Strong | ParetoDominance::Weak => {
                        props[first_ind_idx].strength += N::one();
                    }
                    ParetoDominance::None => {
                        props[second_ind_idx].strength += N::one();
                    }
                    _ => {}
                }
            }
        }
    }

    fn set_raw_fitness(&mut self) {
        let objs = self.popul_mp.definitions().objs();
        let mp = &self.arch_u_popul.mp;
        let props = &mut self.arch_u_popul.props;
        for (first_ind_idx, first_ind) in mp.iter().enumerate() {
            for (second_ind_idx, second_ind) in mp.iter().enumerate().skip(first_ind_idx) {
                match verify_pareto_dominance(first_ind.objs(), second_ind.objs(), |idx, a, b| {
                    match objs[idx].obj_direction() {
                        ObjDirection::Max => b.partial_cmp(a).unwrap(),
                        ObjDirection::Min => a.partial_cmp(b).unwrap(),
                    }
                }) {
                    ParetoDominance::Strong | ParetoDominance::Weak => {
                        props[second_ind_idx].raw_fitness += props[first_ind_idx].strength;
                    }
                    ParetoDominance::None => {
                        props[first_ind_idx].raw_fitness += props[second_ind_idx].strength;
                    }
                    _ => {}
                }
            }
        }
    }

    fn best_popul_result(&self) -> usize {
        let objs = self.popul_mp.definitions().objs();
        let results = self.popul_mp.results();
        let mut best_idx = 0;
        for idx in 1..results.len() {
            let best = results.get(best_idx);
            let current = results.get(idx);
            match verify_pareto_dominance(current.objs(), best.objs(), |idx, a, b| {
                match objs[idx].obj_direction() {
                    ObjDirection::Max => b.partial_cmp(a).unwrap(),
                    ObjDirection::Min => a.partial_cmp(b).unwrap(),
                }
            }) {
                ParetoDominance::Strong | ParetoDominance::Weak => {
                    if current
                        .hard_cstrs()
                        .iter()
                        .all(|x| x.has_some_weight() == false)
                    {
                        best_idx = idx;
                    }
                }
                _ => {}
            };
        }
        best_idx
    }
}

impl<'a, C, CR, MA, MU, N, S, V> Solver<N, Mphos<'a, C, N, V>, V>
    for Spea2<'a, C, CR, MA, MU, N, S, V>
where
    C: Send + Sync,
    CR: Crossover<N, V>,
    MA: MatingSelection<N, V>,
    MU: Mutation<C, N, V>,
    N: AddAssign + Debug + Float + Send + Sync,
    S: SurvivorSelection,
    V: Copy + Debug + Send + Sync,
{
    fn best_result_objs(&self) -> &[N] {
        self.popul_mp.results().best_objs().unwrap()
    }

    fn best_result_objs_avg(&self) -> N {
        *self.popul_mp.results().best().unwrap().objs_avg()
    }

    fn do_work_after(&mut self) {
        let filling_num = self.popul_mp.definitions().results_num();
        self.gap.mating_selection.mating_selection(
            &mut self.arch_mp,
            self.popul_mp.results_mut(),
            filling_num,
        );
        self.gap.crossover.crossover(self.popul_mp.results_mut());
        self.gap.mutation.mutation(&mut self.popul_mp);
        self.gap.survivor_selection.survivor_selection();
    }
    fn do_work_before(&mut self) {
        {
            let (definitions, results) = self.popul_mp.parts_mut();
            MphosEvaluators::evaluate_results_without_reasons(definitions, &mut self.arch_mp);
            MphosEvaluators::evaluate_results_without_reasons(definitions, results);
        }
        self.arch_u_popul.mp.append(&self.arch_mp);
        self.arch_u_popul.mp.append(&self.popul_mp.results_mut());
        self.assign_fitness();
        self.environment_selection();
        self.arch_u_popul.mp.clear();
        let best_idx = self.best_popul_result();
        self.popul_mp.results_mut().set_best(best_idx);
    }

    fn into_result(self) -> Mphos<'a, C, N, V> {
        self.popul_mp
    }
}