use std::{convert::Infallible, marker::PhantomData};

use eoa_lib::{fitness::FitnessFunction, initializer::Initializer, perturbation::PerturbationOperator};
use itertools::Itertools;
use nalgebra::{allocator::Allocator, distance, Const, DefaultAllocator, Dim, Dyn, OMatrix, OVector, Point, U1};
use rand::{seq::SliceRandom, Rng, RngCore};

#[derive(PartialEq, Clone, Debug)]
pub struct TSPCity {
    point: Point<f64, 2>
}

#[derive(PartialEq, Clone, Debug)]
pub struct NodePermutation<D: Dim>
where
    DefaultAllocator: Allocator<D>
{
    permutation: OVector<usize, D>
}

/// An instance of TSP, a fully connected graph
/// with cities that connect to each other.
/// The D parameter represents the number of cities.
#[derive(PartialEq, Clone, Debug)]
pub struct TSPInstance<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>
{
    cities: Vec<TSPCity>,
    distances: OMatrix<f64, D, D>
}

impl TSPInstance<Dyn>
where
{
    pub fn new_dyn(cities: Vec<(f64, f64)>) -> Self {
        let dim = Dyn(cities.len());

        let cities = OMatrix::<f64, Dyn, Const<2>>::from_fn_generic(dim, Const::<2>, |i, j| if j == 0 { cities[i].0 } else { cities[i].1 });
        TSPInstance::new(cities)
    }
}

impl<const D: usize> TSPInstance<Const<D>>
where
{
    pub fn new_const(cities: Vec<(f64, f64)>) -> Self {
        let cities = OMatrix::<f64, Const<D>, Const<2>>::from_fn(|i, j| if j == 0 { cities[i].0 } else { cities[i].1 });
        TSPInstance::new(cities)
    }
}

impl<D> TSPInstance<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
    DefaultAllocator: Allocator<D>,
    DefaultAllocator: Allocator<D, Const<2>>,
{
    pub fn new(cities: OMatrix<f64, D, Const<2>>) -> Self {
        let dim = cities.shape_generic().0;

        let cities = cities.column_iter()
                .map(|position|
                     TSPCity { point: Point::<f64, 2>::new(position[0], position[1])  }
                )
                .collect::<Vec<_>>();

        let distances = OMatrix::from_fn_generic(
            dim,
            dim,
            |i, j| distance(&cities[i].point, &cities[j].point)
        );

        Self {
            cities,
            distances
        }
    }
}

impl<D> TSPInstance<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
    DefaultAllocator: Allocator<D>,
{
    pub fn verify_solution(&self, solution: &NodePermutation<D>) -> bool {
        let mut seen_vertices = OVector::from_element_generic(
            solution.permutation.shape_generic().0,
            solution.permutation.shape_generic().1,
            false
        );

        for &vertex in solution.permutation.iter() {
            // This vertex index is out of bounds
            if vertex >= self.cities.len() {
                return false;
            }

            // A node is repeating
            if seen_vertices[vertex] {
                return false;
            }

            seen_vertices[vertex] = true;
        }

        true
    }

    pub fn solution_cost(&self, solution: &NodePermutation<D>) -> f64 {
        solution.permutation
            .iter()
            .circular_tuple_windows()
            .map(|(&node1, &node2): (&usize, &usize)| self.distances[(node1, node2)])
            .sum()
    }
}

impl<D> FitnessFunction for TSPInstance<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
    DefaultAllocator: Allocator<D>,
{
    type In = NodePermutation<D>;
    type Out = f64;
    type Err = Infallible;

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        Ok(self.solution_cost(inp))
    }
}

pub struct TSPRandomInitializer<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
{
    _phantom: PhantomData<D>
}

impl<D> Initializer<D, NodePermutation<D>> for TSPRandomInitializer<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
    DefaultAllocator: Allocator<D>,
{
    fn initialize_single(&self, size: D, rng: &mut dyn RngCore) -> NodePermutation<D> {
        let len = size.value();
        let mut indices = OVector::<usize, D>::from_iterator_generic(size, U1, 0..len);
        indices.as_mut_slice().shuffle(rng);

        NodePermutation { permutation: indices }
    }
}

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

impl<D> PerturbationOperator for SwapPerturbation<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D, D>,
    DefaultAllocator: Allocator<D>,
{
    type Chromosome = NodePermutation<D>;

    fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore) {
        let first = rng.random_range(0..=chromosome.permutation.len());
        let second = rng.random_range(0..=chromosome.permutation.len());

        (
            chromosome.permutation[first],
            chromosome.permutation[second]
        ) = (
            chromosome.permutation[second],
            chromosome.permutation[first]
        );
    }
}