use std::marker::PhantomData;

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

use crate::{binary_string::BinaryString, pairing::ParentPairing, replacement::{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 (
                chromosome1,
                chromosome2
            ) = (&parent1.chromosome, &parent2.chromosome);

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