use std::marker::PhantomData;

use nalgebra::{allocator::Allocator, DefaultAllocator, Dim, OVector, Scalar, U1};
use rand::{seq::IteratorRandom, Rng, RngCore};

use crate::{binary_string::BinaryString, pairing::ParentPairing, population::{EvaluatedPopulation, Population}};

pub trait Crossover<const D: usize> {
    type Chromosome;
    type Out;

    fn crossover(
        &self,
        parents: &EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = ParentPairing<D>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome>;
}

pub struct BinaryOnePointCrossover<D: Dim, TOutput> {
    _phantom1: PhantomData<D>,
    _phantom2: PhantomData<TOutput>
}

impl<D: Dim, TOutput> BinaryOnePointCrossover<D, TOutput>
where
    DefaultAllocator: Allocator<D>
{
    pub fn new() -> Self {
        Self {
            _phantom1: PhantomData,
            _phantom2: PhantomData,
        }
    }

    fn find_cross_point(&self, chromosome: &BinaryString<D>, rng: &mut dyn RngCore) -> usize {
        let (min, max) = (0, chromosome.vec.len());
        rng.random_range(min..max)
    }
}

// TODO: make common functions for ovector that will be used from both BinaryOnePointCrossover and OVectorOnePointCrossover
// for not repeating the code.

impl<D, TOutput> Crossover<2> for BinaryOnePointCrossover<D, TOutput>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    type Chromosome = BinaryString<D>;
    type Out = TOutput;

    fn crossover(
        &self,
        population: &EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {

        let chromosome = &population.population[0].chromosome.vec;
        let len = population.population[0].chromosome.vec.len();
        let indices = OVector::<usize, D>::from_iterator_generic(chromosome.shape_generic().0, U1, 0..len);

        let mut offsprings = Vec::new();
        for pair in pairs {
            let (
                parent1,
                parent2
            ) = (&population.population[pair.x], &population.population[pair.y]);

            let (
                mut chromosome1,
                mut chromosome2
            ) = (&parent1.chromosome, &parent2.chromosome);

            if rng.random_bool(0.5) {
                (chromosome1, chromosome2) = (chromosome2, chromosome1)
            }

            let cross_point = self.find_cross_point(&population.population[0].chromosome, rng);

            offsprings.push(BinaryString::from_ovector(
                chromosome1.vec.zip_zip_map(
                    &chromosome2.vec,
                    &indices,
                    |first, second, i| if i <= cross_point { first } else { second }
                )));
        }

        Population::from_vec(offsprings)
    }
}

pub struct BinaryTwoPointCrossover<D: Dim, TOutput> {
    _phantom1: PhantomData<D>,
    _phantom2: PhantomData<TOutput>
}

impl<D: Dim, TOutput> BinaryTwoPointCrossover<D, TOutput>
where
    DefaultAllocator: Allocator<D>
{
    pub fn new() -> Self {
        Self {
            _phantom1: PhantomData,
            _phantom2: PhantomData,
        }
    }

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

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

impl<D, TOutput> Crossover<2> for BinaryTwoPointCrossover<D, TOutput>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    type Chromosome = BinaryString<D>;
    type Out = TOutput;

    fn crossover(
        &self,
        population: &EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {

        let chromosome = &population.population[0].chromosome.vec;
        let len = population.population[0].chromosome.vec.len();
        let indices = OVector::<usize, D>::from_iterator_generic(chromosome.shape_generic().0, U1, 0..len);

        let mut offsprings = Vec::new();
        for pair in pairs {
            let (
                parent1,
                parent2
            ) = (&population.population[pair.x], &population.population[pair.y]);

            let (
                mut chromosome1,
                mut chromosome2
            ) = (&parent1.chromosome, &parent2.chromosome);

            if rng.random_bool(0.5) {
                (chromosome1, chromosome2) = (chromosome2, chromosome1)
            }

            let cross_point = self.find_cross_points(&population.population[0].chromosome, rng);

            offsprings.push(BinaryString::from_ovector(
                chromosome1.vec.zip_zip_map(
                    &chromosome2.vec,
                    &indices,
                    |first, second, i| if i <= cross_point[0] {
                        first
                    } else if i < cross_point[1] {
                        second
                    } else {
                        first
                    }
                )));
        }

        Population::from_vec(offsprings)
    }
}

pub struct BinaryNPointCrossover<const N: usize, D: Dim, TOutput> {
    _phantom1: PhantomData<D>,
    _phantom2: PhantomData<TOutput>,
}

impl<const N: usize, D: Dim, TOutput> BinaryNPointCrossover<N, D, TOutput>
where
    DefaultAllocator: Allocator<D>
{
    pub fn new() -> Self {
        Self {
            _phantom1: PhantomData,
            _phantom2: PhantomData,
        }
    }

    fn find_cross_points(&self, chromosome: &BinaryString<D>, rng: &mut dyn RngCore) -> [usize; N] {
        let mut res = [0; N];
        (0..chromosome.vec.len()).choose_multiple_fill(rng, &mut res);
        res.as_mut_slice().sort_unstable();

        res
    }
}

impl<const N: usize, D, TOutput> Crossover<2> for BinaryNPointCrossover<N, D, TOutput>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    type Chromosome = BinaryString<D>;
    type Out = TOutput;

    fn crossover(
        &self,
        population: &EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {

        let chromosome = &population.population[0].chromosome;
        let len = chromosome.vec.len();

        let mut offsprings = Vec::new();
        for pair in pairs {
            let (
                parent1,
                parent2
            ) = (&population.population[pair.x], &population.population[pair.y]);

            let (
                mut chromosome1,
                mut chromosome2
            ) = (&parent1.chromosome.vec, &parent2.chromosome.vec);

            if rng.random_bool(0.5) {
                (chromosome1, chromosome2) = (chromosome2, chromosome1)
            }

            let parents = [chromosome1, chromosome2];

            let cross_points = self.find_cross_points(chromosome, rng);

            let mut offspring =
                OVector::zeros_generic(chromosome1.shape_generic().0, U1);

            let mut start = 0;
            let mut parent = 0;
            for cross_point in cross_points {
                for i in start..cross_point {
                    offspring[i] = parents[parent][i];
                }

                parent = 1 - parent; // switch parent
                start = cross_point;
            }

            for i in start..len {
                offspring[i] = parents[parent][i];
            }

            offsprings.push(BinaryString::from_ovector(offspring));
        }

        Population::from_vec(offsprings)
    }
}

pub struct OVectorOnePointCrossover<D: Dim, T: Scalar, TOutput> {
    _phantom1: PhantomData<D>,
    _phantom2: PhantomData<TOutput>,
    _phantom3: PhantomData<T>
}

impl<D: Dim, T, TOutput> OVectorOnePointCrossover<D, T, TOutput>
where
    T: Scalar,
    DefaultAllocator: Allocator<D>
{
    pub fn new() -> Self {
        Self {
            _phantom1: PhantomData,
            _phantom2: PhantomData,
            _phantom3: PhantomData,
        }
    }

    fn find_cross_point(&self, chromosome: &OVector<T, D>, rng: &mut dyn RngCore) -> usize {
        let (min, max) = (0, chromosome.len() - 1);
        rng.random_range(min..max)
    }
}

impl<D, T, TOutput> Crossover<2> for OVectorOnePointCrossover<D, T, TOutput>
where
    T: Scalar,
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    type Chromosome = OVector<T, D>;
    type Out = TOutput;

    fn crossover(
        &self,
        population: &EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {

        let chromosome = &population.population[0].chromosome;
        let len = population.population[0].chromosome.len();
        let indices = OVector::<usize, D>::from_iterator_generic(chromosome.shape_generic().0, U1, 0..len);

        let mut offsprings = Vec::new();
        for pair in pairs {
            let (
                parent1,
                parent2
            ) = (&population.population[pair.x], &population.population[pair.y]);

            let (
                chromosome1,
                chromosome2
            ) = (&parent1.chromosome, &parent2.chromosome);

            let cross_point = self.find_cross_point(&population.population[0].chromosome, rng);

            offsprings.push(
                chromosome1.zip_zip_map(
                    &chromosome2,
                    &indices,
                    |first, second, i| if i <= cross_point { first } else { second }
                ));
        }

        Population::from_vec(offsprings)
    }
}

pub struct ArithmeticCrossover<D: Dim, TOutput> {
    _phantom1: PhantomData<D>,
    _phantom2: PhantomData<TOutput>,
}

impl<D, TOutput> ArithmeticCrossover<D, TOutput> {
    pub fn new() -> Self {
        Self {
            _phantom1: PhantomData,
            _phantom2: PhantomData,
        }
    }
}

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

    fn crossover(
        &self,
        population: &EvaluatedPopulation<Self::Chromosome, Self::Out>,
        pairs: impl Iterator<Item = ParentPairing<2>>,
        rng: &mut dyn RngCore
    ) -> Population<Self::Chromosome> {
        let mut offsprings = Vec::new();

        for pair in pairs {
            let (
                parent1,
                parent2
            ) = (&population.population[pair.x], &population.population[pair.y]);

            let (
                chromosome1,
                chromosome2
            ) = (&parent1.chromosome, &parent2.chromosome);

            offsprings.push(
                chromosome1.zip_map(
                    &chromosome2,
                    |first, second| {
                        let alpha = rng.random_range(0.0..=1.0);
                        alpha * first + (1.0 - alpha) * second
                    }
                ));

            offsprings.push(
                chromosome1.zip_map(
                    &chromosome2,
                    |first, second| {
                        let alpha = rng.random_range(0.0..=1.0);
                        alpha * second + (1.0 - alpha) * first
                    }
                ));
        }

        Population::from_vec(offsprings)
    }
}