use std::{cmp::Ordering, marker::PhantomData};
use nalgebra::{allocator::Allocator, Const, DefaultAllocator, Dim, OVector, U1, U2};
use rand::{prelude::SliceRandom, seq::IteratorRandom, Rng, RngCore};
use eoa_lib::initializer::Initializer;
use crate::{graph::{Graph, minimal_spanning_tree_kruskal, GenericGraph}, tsp::{NodePermutation, TSPCity, TSPEdge, 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 _ in 0..count {
            population.push(self.initialize_single(size, rng))
        }

        population
    }
}

pub struct MinimumSpanningTreeInitializer<'a, D: Dim>
where
    DefaultAllocator: nalgebra::allocator::Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, D>
{
    instance: &'a TSPInstance<D>,
    minimal_spanning_tree: GenericGraph<TSPCity, TSPEdge>,
    _phantom: PhantomData<D>,
}

impl<'a, D: Dim> MinimumSpanningTreeInitializer<'a, D>
where
    DefaultAllocator: nalgebra::allocator::Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, U2>,
{
    pub fn new(instance: &'a TSPInstance<D>) -> Self {
        let tsp_graph = instance.clone().to_graph();
        let minimal_spanning_tree =
            minimal_spanning_tree_kruskal(
                &tsp_graph,
                None,
                |_| true
            );

        assert_eq!(minimal_spanning_tree.components_count(), 1);
        let minimal_spanning_tree = tsp_graph
            .filter_edges(|i, _|
                          minimal_spanning_tree.edges.contains(&i));

        Self {
            instance,
            minimal_spanning_tree,
            _phantom: PhantomData
        }
    }

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

        let mut current_node = node;
        for i in 0..size.value()-1 {
            individual[i] = current_node;
            used_nodes[current_node] = true;

            let usable_neighbors = self.minimal_spanning_tree
                .neighbor_idxs(current_node)
                .unwrap()
                .filter(|&neighbor| !used_nodes[neighbor]);

            let chosen = usable_neighbors.choose(rng);

            current_node = if let Some(next_node) = chosen {
                next_node
            } else {
                // Choose closest node
                (0..size.value())
                    .filter(|&node| !used_nodes[node])
                    .min_by(|&next_node1, &next_node2|
                            self.instance.distances(current_node, next_node1)
                            .partial_cmp(&self.instance.distances(current_node, next_node2))
                            .unwrap_or(Ordering::Less))
                    // There has to be unused node since we haven't yet used size.value() nodes.
                    .unwrap()
            }
        }

        // The last node
        individual[size.value() - 1] = current_node;

        NodePermutation { permutation: individual }
    }
}

impl<'a, D: Dim> Initializer<D, NodePermutation<D>> for MinimumSpanningTreeInitializer<'a, D>
where
    DefaultAllocator: nalgebra::allocator::Allocator<D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, D>,
    DefaultAllocator: nalgebra::allocator::Allocator<D, U2>,
{
    fn initialize_single(&self, size: D, rng: &mut dyn RngCore) -> NodePermutation<D> {
        let starting_node = rng.random_range(0..size.value());
        self.initialize_from(starting_node, size, 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 every node
        for i in 0..size.value() {
            population.push(self.initialize_from(
                i,
                size,
                rng
            ));

            count -= 1;

            if count == 0 {
                return population;
            }
        }

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

        population
    }
}