use std::marker::PhantomData;
use nalgebra::{allocator::Allocator, Const, DefaultAllocator, Dim, OMatrix, OVector, U1};
use rand::{prelude::IteratorRandom, Rng, RngCore};
use eoa_lib::replacement::Population;
use itertools::Itertools;
use eoa_lib::crossover::Crossover;
use crate::tsp::NodePermutation;

pub struct NoCrossover<D> {
    _phantom: PhantomData<D>
}

impl<D> NoCrossover<D> {
    pub fn new() -> Self {
        Self { _phantom: PhantomData }
    }
}

impl<D> Crossover<2> for NoCrossover<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    type Chromosome = NodePermutation<D>;
    type Out = f64;

    fn crossover(
        &self,
        parents: &eoa_lib::replacement::EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = eoa_lib::pairing::ParentPairing<2>>,
        _: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {
        let mut offsprings = vec![];
        for pair in pairs {
            offsprings.push(parents.population[pair[0]].chromosome.clone());
            offsprings.push(parents.population[pair[1]].chromosome.clone());
        }

        Population::from_vec(offsprings)
    }
}

pub struct EdgeRecombinationCrossover<D> {
    _phantom: PhantomData<D>
}

impl<D> EdgeRecombinationCrossover<D> {
    pub fn new() -> Self {
        Self { _phantom: PhantomData }
    }
}

impl<D> Crossover<2> for EdgeRecombinationCrossover<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
    DefaultAllocator: Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, Const<4>>
{
    type Chromosome = NodePermutation<D>;
    type Out = f64;

    fn crossover(
        &self,
        parents: &eoa_lib::replacement::EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = eoa_lib::pairing::ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> eoa_lib::replacement::Population<Self::Chromosome> {
        let mut offsprings = vec![];

        let permutation = &parents.population[0].chromosome.permutation;
        let len = permutation.len();
        let mut adjacency_lists = OMatrix::from_element_generic(
            permutation.shape_generic().0,
            Const::<4>,
            None);
        let mut used_nodes = OVector::from_element_generic(
            permutation.shape_generic().0,
            Const::<1>,
            false
        );

        let mut neighbors_count = OVector::from_element_generic(
            permutation.shape_generic().0,
            Const::<1>,
            2usize
        );

        for pair in pairs {
            let parent1 = &parents.population[pair.x].chromosome;
            let parent2 = &parents.population[pair.y].chromosome;

            used_nodes.apply(|n| *n = false);

            // 1. Populate adjacency lists
            for (&c1, &n, &c2) in parent1.permutation.iter().circular_tuple_windows() {
                adjacency_lists[(n, 0)] = Some(c1);
                adjacency_lists[(n, 1)] = Some(c2);
                neighbors_count[n] = 2;
            }

            for (&c1, &n, &c2) in parent2.permutation.iter().circular_tuple_windows() {
                // Not duplicit?
                if adjacency_lists[(n, 0)].unwrap() != c1 && adjacency_lists[(n, 1)].unwrap() != c1 {
                    neighbors_count[n] += 1;
                    adjacency_lists[(n, 2)] = Some(c1);
                } else { // Duplicit
                    adjacency_lists[(n, 2)] = None;
                }

                // Not duplicit
                if adjacency_lists[(n, 0)].unwrap() != c2 && adjacency_lists[(n, 1)].unwrap() != c2 {
                    neighbors_count[n] += 1;
                    adjacency_lists[(n, 3)] = Some(c2);
                } else { // Duplicit
                    adjacency_lists[(n, 3)] = None;
                }
            }

            let chosen_parent = if rng.random_bool(0.5) {
                &parent1
            } else {
                &parent2
            };

            let mut offspring = OVector::from_element_generic(permutation.shape_generic().0, Const::<1>, 0);

            let mut current_node = chosen_parent.permutation[0];

            for i in 0..len-1 {
                offspring[i] = current_node;
                used_nodes[current_node] = true;

                for neighbor in adjacency_lists.row(current_node) {
                    if let Some(neighbor) = neighbor {
                        neighbors_count[*neighbor] -= 1;
                    }
                }

                let min_neighbors = adjacency_lists.row(current_node)
                    .iter()
                    .flatten()
                    .filter(|&&neighbor| !used_nodes[neighbor])
                    .map(|&neighbor| neighbors_count[neighbor])
                    .min();

                let neighbor = if let Some(min_neighbors) = min_neighbors {
                    adjacency_lists.row(current_node)
                        .iter()
                        .flatten()
                        .copied()
                        .filter(|&neighbor| !used_nodes[neighbor] && neighbors_count[neighbor] == min_neighbors)
                        .choose(rng)
                } else {
                    None
                };

                current_node = if let Some(neighbor) = neighbor {
                    neighbor
                } else {
                    (0..len).filter(|&node| !used_nodes[node])
                    .choose(rng)
                    .unwrap()
                };
            }

            offspring[len - 1] = current_node;

            offsprings.push(NodePermutation { permutation: offspring });
        }

        Population::from_vec(offsprings)
    }
}

pub struct CycleCrossover<D: Dim> {
    _phantom: PhantomData<D>
}

impl<D: Dim> CycleCrossover<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    pub fn new() -> Self {
        Self { _phantom: PhantomData }
    }

    fn perform_crossover(
        &self,
        parent1: &OVector<usize, D>,
        parent2: &OVector<usize, D>,
        city_positions: &mut [usize]
    ) -> NodePermutation<D> {
        let mut offspring = parent2.clone();

        for (i, &city) in parent1.iter().enumerate() {
            city_positions[city] = i;
        }

        let mut i = 0;
        let mut first = true;
        while i != 0 || first {
            first = false;

            let city = parent1[i];

            offspring[i] = city;

            let city = parent2[i];
            i = city_positions[city];
        }

        NodePermutation { permutation: offspring }
    }
}

impl<D: Dim> Crossover<2> for CycleCrossover<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    type Chromosome = NodePermutation<D>;
    type Out = f64;

    fn crossover(
        &self,
        parents: &eoa_lib::replacement::EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = eoa_lib::pairing::ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {
        let mut offsprings = vec![];

        let permutation = &parents.population[0].chromosome.permutation;
        let mut city_positions =
            OVector::zeros_generic(permutation.shape_generic().0, U1);

        for pair in pairs {
            let parent1 = &parents.population[pair.x].chromosome;
            let parent2 = &parents.population[pair.y].chromosome;

            let (perm1, perm2) = (
                &parent1.permutation,
                &parent2.permutation
            );

            offsprings.push(self.perform_crossover(perm1, perm2, city_positions.as_mut_slice()));
            offsprings.push(self.perform_crossover(perm2, perm1, city_positions.as_mut_slice()));
        }

        Population::from_vec(offsprings)
    }
}

pub struct PartiallyMappedCrossover<D: Dim> {
    _phantom: PhantomData<D>
}

impl<D: Dim> PartiallyMappedCrossover<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    pub fn new() -> Self {
        Self { _phantom: PhantomData }
    }

    fn find_cross_points(&self, chromosome: &NodePermutation<D>, rng: &mut dyn RngCore) -> [usize; 2] {
        let (min, max) = (0, chromosome.permutation.len());
        let first = rng.random_range(min..max);
        let second = rng.random_range(min..max);

        [ first.min(second), first.max(second) ]
    }

    fn perform_crossover(
        &self,
        parent1: &OVector<usize, D>,
        parent2: &OVector<usize, D>,
        crossover_points: &[usize; 2],
        city_positions: &mut [usize]
    ) -> NodePermutation<D> {
        let mut offspring = parent1.clone();

        for (i, &city) in parent1.iter().enumerate() {
            city_positions[city] = i;
        }

        for i in crossover_points[0]..crossover_points[1] {
            let city = parent2[i];

            offspring.swap_rows(
                i,
                city_positions[city]
            );
        }

        NodePermutation { permutation: offspring }
    }
}

impl<D: Dim> Crossover<2> for PartiallyMappedCrossover<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    type Chromosome = NodePermutation<D>;
    type Out = f64;

    fn crossover(
        &self,
        parents: &eoa_lib::replacement::EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = eoa_lib::pairing::ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {
        let mut offsprings = vec![];

        let permutation = &parents.population[0].chromosome.permutation;
        let mut city_positions =
            OVector::zeros_generic(permutation.shape_generic().0, U1);

        for pair in pairs {
            let parent1 = &parents.population[pair.x].chromosome;
            let parent2 = &parents.population[pair.y].chromosome;

            let (perm1, perm2) = (
                &parent1.permutation,
                &parent2.permutation
            );

            let cross_points = self.find_cross_points(parent1, rng);
            offsprings.push(self.perform_crossover(perm1, perm2, &cross_points, city_positions.as_mut_slice()));
            offsprings.push(self.perform_crossover(perm2, perm1, &cross_points, city_positions.as_mut_slice()));
        }

        Population::from_vec(offsprings)
    }
}