use crate::colony::Colony;
use crate::cycle::brent_cycle;
use crate::petgraph::graph::NodeIndex;
use crate::petgraph::EdgeType;
use crate::Objective;
use rand::distributions::{Distribution, Uniform};
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use std::cmp::PartialEq;
use std::fmt::{Debug, Display};
use std::marker::Send;

#[derive(Debug)]
pub struct Ant {
    pub number: usize,
    pub visited_nodes: Vec<usize>,
    pub visited_edges: Vec<usize>,
    pub edge_weight: f64,
    obj: Objective,
    rng: SmallRng,
    cycling: bool,
}

impl Ant {
    pub fn from_node(number: usize, node_idx: usize, seed: u64, obj: Objective) -> Ant {
        Ant {
            number,
            visited_nodes: vec![node_idx],
            visited_edges: Vec::new(),
            edge_weight: 0.0,
            obj,
            rng: SmallRng::seed_from_u64(seed),
            cycling: false,
        }
    }

    pub fn move_prob<
        N: 'static + Clone + Display + Debug + PartialEq + Send,
        Ty: 'static + Clone + EdgeType + Send,
    >(
        &self,
        colony: &Colony<N, Ty>,
        from: usize,
        to: usize,
    ) -> f64 {
        let numerator = colony.unnormalized_probability(from, to);
        let mut denominator = 0.0;

        let neighbors = colony.pheromone_graph.neighbors(NodeIndex::new(from));
        // Compute the denominator from the neightboring nodes only
        for nodeidx in neighbors {
            if !self.visited_nodes.contains(&nodeidx.index()) || self.cycling {
                denominator += colony.unnormalized_probability(from, nodeidx.index());
            }
        }
        numerator / denominator
    }

    /// Moves an ant to the next node. Returns false when the objective is satisfied.
    pub fn move_ant<
        N: 'static + Clone + Display + Debug + PartialEq + Send,
        Ty: 'static + Clone + EdgeType + Send,
    >(
        &mut self,
        colony: &Colony<N, Ty>,
    ) -> bool {
        match self.obj {
            Objective::VisitAllNodes => {
                if self.visited_nodes.len() == colony.pheromone_graph.node_count() {
                    // If we have a complete graph, then this condition exists
                    debug!("Ant#{} has visited all nodes", self.number);
                    return false;
                }
            }
            Objective::VisitNNodes(n) => {
                if self.visited_nodes.len() == n {
                    // If we have a complete graph, then this condition exists
                    debug!("Ant#{} has visited {} nodes", self.number, n);
                    return false;
                }
            }
            Objective::ReachNodeIdx(idx) => {
                if self.visited_nodes.contains(&idx) {
                    // If we have a complete graph, then this condition exists
                    debug!(
                        "Ant#{} has reached destination node {} (idx: {})",
                        self.number,
                        colony.viz_graph[NodeIndex::new(idx)],
                        idx
                    );
                    return false;
                }
            }
            _ => {}
        }

        // Move to the next node which has the highest probability based on the connected nodes.
        let current_node = *self.visited_nodes.last().unwrap();
        let neighbors = colony
            .pheromone_graph
            .neighbors(NodeIndex::new(current_node));

        let mut next_node = 0;

        if self.rng.gen::<f64>() < colony.opt.wander {
            // Take a random step regardless of the pheromones
            // NOTE: We always need to choose from the neighbors instead of guessing the next node
            // or we may try to create a link when there is none in an incomplete graph.
            debug!("Ant#{} wandering", self.number);
            // Let's check that there are some neighbors which have not been visited.
            let num_neighbors = neighbors.clone().count();
            let mut all_neighbors_visited = true;
            let mut available_neighbors = Vec::new();
            for neighbor in neighbors.clone() {
                if !self.visited_nodes.contains(&neighbor.index()) {
                    all_neighbors_visited = false;
                    // And set the next node to this.
                    available_neighbors.push(neighbor);
                }
            }
            if all_neighbors_visited {
                match self.obj {
                    Objective::ReachDeadEnd => {
                        debug!(
                            "Ant#{} has hit a dead end after {} nodes",
                            self.number,
                            self.visited_nodes.len()
                        );
                        // Let's normalize the weight over the number of nodes
                        self.edge_weight /= self.visited_nodes.len() as f64;
                        return false;
                    }
                    _ => {
                        // If there are no neighbors, we're done for!
                        if neighbors.clone().count() == 0 {
                            warn!(
                                "Ant#{} has hit a dead end and there are no neighbors",
                                self.number,
                            );
                            return false;
                        }
                        // Return to the path with the highest pheromone trail
                        // Go to the node with the highest pheromone levels.
                        let mut max_prob = 0.0;
                        // Mark this ant as currently cycling to compute the probability all of
                        // nodes and not simply those which have not been visited.
                        self.cycling = true;
                        // Find neighboring nodes
                        for nodeidx in neighbors {
                            let next_prob = self.move_prob(&colony, current_node, nodeidx.index());
                            debug!(
                                "Ant#{} - {} -> {}: {}",
                                self.number,
                                colony.viz_graph[NodeIndex::new(current_node)],
                                colony.viz_graph[nodeidx],
                                next_prob
                            );
                            if next_prob > max_prob {
                                max_prob = next_prob;
                                next_node = nodeidx.index();
                            }
                        }
                        debug!(
                            "Ant#{} reached a dead end, going backward to {} with probability {}",
                            self.number,
                            colony.viz_graph[NodeIndex::new(next_node)],
                            max_prob
                        );
                        self.cycling = false; // Node selected, we're done cycling
                    }
                }
            } else {
                // Randomly chose a neighbor from those available.
                let next_node_idx =
                    Uniform::from(0..available_neighbors.len()).sample(&mut self.rng);
                next_node = available_neighbors[next_node_idx].index();
                debug!(
                    "randomly chose {} from {} neighbors, currently at {}",
                    colony.viz_graph[NodeIndex::new(next_node)],
                    num_neighbors,
                    colony.viz_graph[NodeIndex::new(current_node)]
                );
            }
        } else {
            // Go to the node with the highest pheromone levels.
            let mut max_prob = 0.0;
            let mut dead_end = true;
            // Find neighboring nodes
            for nodeidx in neighbors {
                if self.visited_nodes.contains(&nodeidx.index()) {
                    continue; // No backtrack
                }
                dead_end = false;
                let next_prob = self.move_prob(&colony, current_node, nodeidx.index());
                debug!(
                    "Ant#{} - {} -> {}: {}",
                    self.number,
                    colony.viz_graph[NodeIndex::new(current_node)],
                    colony.viz_graph[nodeidx],
                    next_prob
                );
                if next_prob > max_prob {
                    max_prob = next_prob;
                    next_node = nodeidx.index();
                }
            }
            if dead_end {
                match self.obj {
                    Objective::ReachDeadEnd => {
                        debug!(
                            "Ant#{} has hit a dead end after {} nodes",
                            self.number,
                            self.visited_nodes.len()
                        );
                        // Let's normalize the weight over the number of nodes
                        self.edge_weight /= self.visited_nodes.len() as f64;
                        return false;
                    }
                    _ => {
                        // If we have an incomplete graph, then this condition exists
                        debug!("Ant#{} reached a dead end, going backward!", self.number);
                        next_node = self.visited_nodes[self.visited_nodes.len() - 2];
                        debug!(
                            "next node will be {} from {}",
                            colony.viz_graph[NodeIndex::new(next_node)],
                            colony.viz_graph[NodeIndex::new(current_node)]
                        );
                    }
                }
            }
        }
        // And move and store the edge used.
        self.visited_nodes.push(next_node);
        match colony
            .pheromone_graph
            .find_edge(NodeIndex::new(current_node), NodeIndex::new(next_node))
        {
            Some(edge_idx) => {
                self.visited_edges.push(edge_idx.index());
                self.edge_weight += colony.viz_graph[edge_idx];
                debug!(
                    "Ant#{} moved to {:?} using edge {:?} {:?}",
                    self.number,
                    colony.viz_graph[NodeIndex::new(next_node)],
                    edge_idx,
                    edge_idx.index()
                );
            }
            None => {
                println!(
                    "Ant#{} cannot find edge from {} to {}!",
                    self.number, current_node, next_node
                );
                panic!("it broke");
            }
        }
        // let edge_idx = colony
        //     .pheromone_graph
        //     .find_edge(NodeIndex::new(current_node), NodeIndex::new(next_node))
        //     .unwrap();

        true
    }

    /// Remove any cycles and updates the cost function
    pub fn remove_cycles<
        N: 'static + Clone + Display + Debug + PartialEq + Send,
        Ty: 'static + Clone + EdgeType + Send,
    >(
        &mut self,
        colony: &Colony<N, Ty>,
    ) {
        let new_nodes = brent_cycle(self.visited_nodes.as_slice());
        if new_nodes.len() < self.visited_nodes.len() {
            // Let's recompute the cost of those edges.
            self.edge_weight = 0.0;
            self.visited_edges = Vec::new();
            let mut it = 0;
            while it < new_nodes.len() - 1 {
                debug!(
                    "getting edge between {} and {}",
                    new_nodes[it],
                    new_nodes[it + 1]
                );
                let edge_idx = colony
                    .pheromone_graph
                    .find_edge(
                        NodeIndex::new(new_nodes[it]),
                        NodeIndex::new(new_nodes[it + 1]),
                    )
                    .unwrap();
                self.visited_edges.push(edge_idx.index());
                self.edge_weight += colony.viz_graph[edge_idx];
                it += 1;
            }

            self.visited_nodes = new_nodes;
        }
    }
}