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)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::convert::Infallible;
    use nalgebra::{SVector, U6};
    use rand::{rngs::StdRng, RngCore, SeedableRng};
    use eoa_lib::{fitness::FitnessFunction, initializer::Initializer, pairing::{AdjacentPairing, Pairing}, replacement::Population};
    use crate::initializers::TSPRandomInitializer;
    use crate::tsp::{NodePermutation, TSPInstance};

    struct MockRng;
    impl RngCore for MockRng {
        fn next_u32(&mut self) -> u32 {
            0
        }

        fn next_u64(&mut self) -> u64 {
            0
        }

        fn fill_bytes(&mut self, _: &mut [u8]) {
            panic!()
        }
    }

    struct ZeroFitness<const LEN: usize>;
    impl<const LEN: usize> FitnessFunction for ZeroFitness<LEN> {
        type In = NodePermutation<Const<LEN>>;
        type Out = f64;
        type Err = Infallible;

        fn fit(self: &Self, _: &Self::In) -> Result<Self::Out, Self::Err> {
            Ok(0.0)
        }
    }

    #[test]
    fn test_edge_recombination_properties() {
        let crossover = EdgeRecombinationCrossover::<Const<10>>::new();
        let initializer = TSPRandomInitializer::<Const<10>>::new();
        let adjacency_pairing = AdjacentPairing::new();

        let mut rng = StdRng::seed_from_u64(0);
        for _ in 0..100 {
            let parents = Population::from_vec(initializer.initialize(Const::<10>, 10, &mut rng));
            let parents = parents.evaluate(&ZeroFitness).unwrap();

            let pairs = adjacency_pairing.pair(&parents, 0..10);
            let result = crossover.crossover(&parents, pairs, &mut rng);

            // Test invariants that should always hold:
            for chromosome in result.into_iter() {
                assert!(TSPInstance::verify_solution(&chromosome));
            }
        }
    }

    #[test]
    fn test_edge_recombination_specific_case() {
        let parent1: Vec<usize> = vec![0, 1, 2, 4, 5, 3];
        let parent2: Vec<usize> = vec![2, 0, 1, 3, 4, 5];

        let parent1 = NodePermutation::<U6> { permutation: SVector::<usize, 6>::from_vec(parent1) };
        let parent2 = NodePermutation::<U6> { permutation: SVector::<usize, 6>::from_vec(parent2) };

        let pairing = SVector::<usize, 2>::new(0, 1);
        let pairings = vec![pairing].into_iter();

        let parents = Population::from_vec(vec![parent1, parent2]).evaluate(&ZeroFitness).unwrap();

        let crossover = EdgeRecombinationCrossover::<U6>::new();

        let offsprings = crossover.crossover(&parents, pairings, &mut MockRng);
        let offspring = offsprings.into_iter().next().unwrap();

        // NOTE: this sort of relies on the implementation of the algorithm (when there are multiple possibilities
        // currently the algorithm always chooses last). It's possible this test will break due to valid changes to the algorithm.
        assert_eq!(vec![0usize, 1, 3, 4, 5, 2], offspring.permutation.into_iter().copied().collect::<Vec<_>>())
    }
}