use std::{cmp::Ordering, marker::PhantomData};
use nalgebra::{allocator::Allocator, DefaultAllocator, Dim, OVector, U1};
use rand::{prelude::SliceRandom, Rng, RngCore};
use eoa_lib::initializer::Initializer;
use crate::tsp::{NodePermutation, TSPInstance};

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

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

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 NearestNeighborInitializer<'a, D: Dim>
where
    DefaultAllocator: nalgebra::allocator::Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, D>
{
    instance: &'a TSPInstance<D>,
    neighbors_order: Vec<Vec<usize>>,
    _phantom: PhantomData<D>,
}

impl<'a, D: Dim> NearestNeighborInitializer<'a, D>
where
    DefaultAllocator: nalgebra::allocator::Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, D>
{
    pub fn new(instance: &'a TSPInstance<D>) -> Self {
        let mut neighbors_order = vec![
            (0..instance.cities.len()).collect::<Vec<_>>();
            instance.cities.len()
        ];

        for node in 0..instance.cities.len() {
            let neighbor_order = &mut neighbors_order[node];
            neighbor_order.sort_by(
                |&neighbor1, &neighbor2| instance.distances(node, neighbor1)
                    .partial_cmp(&instance.distances(node, neighbor2))
                    .unwrap_or(Ordering::Less)
            );
        }

        Self {
            neighbors_order,
            instance,
            _phantom: PhantomData
        }
    }

    fn initialize_from(
        &self,
        size: D,
        index: usize,
        p_choose_second: f64,
        rng: &mut dyn RngCore
    ) -> NodePermutation<D> {
        let mut used_neighbors = vec![false; size.value()];
        let mut individual =
            OVector::zeros_generic(size, U1);

        let mut current_node = index;
        individual[0] = current_node;
        used_neighbors[current_node] = true;

        for i in 1..size.value() {
            let neighbor_order = &self.neighbors_order[current_node];

            let mut next_node = None;
            for &neighbor in neighbor_order {
                if !used_neighbors[neighbor] {
                    next_node = Some(neighbor);

                    // Do not choose third, ...
                    if next_node.is_some() {
                        break;
                    }

                    // Possibility to choose second nearest neighbor
                    if !rng.random_bool(p_choose_second) {
                        break;
                    }
                }
            }

            current_node = next_node.unwrap();
            used_neighbors[current_node] = true;
            individual[i] = current_node;
        }

        NodePermutation { permutation: individual }
    }
}

impl<'a, D: Dim> Initializer<D, NodePermutation<D>> for NearestNeighborInitializer<'a, D>
where
    DefaultAllocator: nalgebra::allocator::Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, D>
{
    fn initialize_single(&self, size: D, rng: &mut dyn RngCore) -> NodePermutation<D> {
        let starting_node = rng.random_range(0..size.value());
        self.initialize_from(size, starting_node, rng.random_range(0.0..=0.7), rng)
    }

    fn initialize(
        &self,
        size: D,
        mut count: usize,
        rng: &mut dyn RngCore
    ) -> Vec<NodePermutation<D>> {
        let mut population = Vec::with_capacity(count);

        // 1. Initialize from all starting nodes with nearest (p_choose_second == 0)
        for i in 0..size.value() {
            population.push(self.initialize_from(
                size,
                i,
                0.0,
                rng
            ));

            count -= 1;

            if count == 0 {
                return population;
            }
        }

        // 2. Initialize from all starting nodes with second nearest (p_choose_second == 1)
        for i in 0..size.value() {
            population.push(self.initialize_from(
                size,
                i,
                1.0,
                rng
            ));

            count -= 1;

            if count == 0 {
                return population;
            }
        }

        // 3. Initialize randomly with random p_choose_second (call initialize_single)
        for i in 0..count {
            population.push(self.initialize_single(size, rng))
        }


        population
    }
}