graco 0.1.3

use crate::ant::Ant;
use crate::io::{export_dot, load_graph};
use crate::petgraph::graph::{Graph, NodeIndex};
use crate::petgraph::prelude::*;
use crate::petgraph::EdgeType;
use crate::Objective;
use std::cmp::PartialEq;
use std::f64::EPSILON;
use std::fmt::{Debug, Display};
use std::marker::Send;
use std::path::PathBuf;
use std::sync::mpsc::channel;
use std::thread::sleep;
use std::time::Duration;
use threadpool::ThreadPool;

#[derive(Clone, Copy, Debug)]
pub struct ColonyOpts<N> {
    /// Verbose mode will output all pheromone graphs at each iteration.
    pub verbose: bool,
    /// Set the initial pheromone in the graph
    pub initphe: f64,
    /// Set the evaporation rate, between 0 and 1, lower faster: 0 is iteration-only pheromones.
    pub evaporation: f64,
    /// Set the maximum number of iterations
    pub iterations: u32,
    /// Set the percentage of ants which need to find the same path to consider the problem solved
    pub threshold: f64,
    /// Density constant Q of the pheromones
    pub density: f64,
    /// Set the relative importance of pheromones for the probabilistic function
    pub alpha: i32,
    /// Set the relative importance of distance for the probabilistic function
    pub beta: i32,
    /// Set the seed of the first ant. Each subsequent ant will have that number incremented
    pub seed: Option<u64>,
    /// Specify the name of the node to reach
    pub visit_node: Option<N>,
    /// Specify the number of nodes to reach
    pub visit_count: Option<usize>,
    /// Set to true if the objective is to reach a dead end
    pub reach_dead_end: bool,
    /// Specify the starting node name
    pub start_node: Option<N>,
    /// Set the probability of not following the pheromone trail
    pub wander: f64,
    /// Number of ants, maxed to the number of nodes in the graph
    pub ants: usize,
    /// Change the number of CPUs to run on (default is max available)
    pub cpus: usize,
}

impl<N: 'static + Clone + Display + Debug + PartialEq + Send> Default for ColonyOpts<N> {
    /// By default, the objective is to visit all nodes
    fn default() -> Self {
        ColonyOpts {
            verbose: false,
            initphe: 1.0,
            evaporation: 0.95,
            iterations: 100,
            threshold: 0.9,
            density: 3.0,
            alpha: 1,
            beta: 2,
            wander: 0.15,
            cpus: 8,
            ants: 50,
            seed: None,
            start_node: None,
            visit_node: None,
            visit_count: None,
            reach_dead_end: false,
        }
    }
}

#[derive(Debug, Clone)]
pub struct Colony<N, Ty: Clone + EdgeType> {
    pub viz_graph: Graph<N, f64, Ty>,
    pub pheromone_graph: Graph<N, f64, Ty>,
    pub opt: ColonyOpts<N>,
    shortest_path: Vec<usize>,
    iteration: u32,
    init_seed: u64,
    obj: Objective,
    start_idx: usize,
}

impl Colony<String, Undirected> {
    pub fn with_input_path(opt: ColonyOpts<String>, input: PathBuf) -> Self {
        let (viz_graph, pheromone_graph) = load_graph(input, opt.initphe);
        Self::from_options(opt, viz_graph, pheromone_graph)
    }
}

impl<
        N: 'static + Clone + Display + Debug + PartialEq + Send,
        Ty: 'static + Clone + EdgeType + Send,
    > Colony<N, Ty>
{
    pub fn from_options(
        opt: ColonyOpts<N>,
        viz_graph: Graph<N, f64, Ty>,
        pheromone_graph: Graph<N, f64, Ty>,
    ) -> Self {
        let seed = match opt.seed {
            Some(seed) => seed,
            None => {
                // Generate a new random number as a seed.
                rand::random::<u64>()
            }
        };

        info!("Seed: {}", seed);

        // Interpret the objective
        let visit_node = opt.visit_node.clone();
        let obj = if let Some(name) = visit_node {
            // Find the destination node index
            let mut idx = 0;
            let mut found = false;
            for node in viz_graph.raw_nodes() {
                if node.weight == name {
                    found = true;
                    break;
                }
                idx += 1;
            }
            if !found {
                panic!("could not find the requested objective node");
            }
            Objective::ReachNodeIdx(idx)
        } else if let Some(count) = opt.visit_count {
            Objective::VisitNNodes(count)
        } else if opt.reach_dead_end {
            Objective::ReachDeadEnd
        } else {
            Objective::VisitAllNodes
        };

        let start_node = opt.start_node.clone();
        let start_idx = if let Some(name) = start_node {
            // Find the destination node index
            let mut idx = 0;
            let mut found = false;
            for node in viz_graph.raw_nodes() {
                if node.weight == name {
                    found = true;
                    break;
                }
                idx += 1;
            }
            if !found {
                panic!("could not find the requested starting node");
            }
            idx
        } else {
            0
        };

        Colony {
            viz_graph,
            pheromone_graph,
            opt,
            shortest_path: Vec::new(),
            iteration: 0,
            init_seed: seed,
            obj: obj,
            start_idx,
        }
    }

    /// The UN-normalized probability between two nodes, i.e. visibility times pheromone.
    pub fn unnormalized_probability(&self, from: usize, to: usize) -> f64 {
        let node1 = NodeIndex::new(from);
        let node2 = NodeIndex::new(to);
        match self.pheromone_graph.find_edge(node1, node2) {
            None => 0.0,
            Some(edge) => {
                // This edge exists.
                let tau = self.pheromone_graph[edge].powi(self.opt.alpha);
                // Find the visibility of this node ("distance").
                let viz = self.viz_graph[edge].powi(self.opt.beta);
                tau * viz
            }
        }
    }

    /// Update the pheromones based on the shortest list of edges.
    pub fn update_pheromones(&mut self, new_path: Vec<usize>) {
        // Export the first pheromone graph if requested
        if self.iteration == 0 && self.opt.verbose {
            self.export();
            export_dot(&self.viz_graph, "dotgraphs/visibility.dot");
        }

        self.shortest_path = new_path.clone();
        // Evaporate all pheromones and increase the pheromones for the used path.
        for edge_no in 0..self.pheromone_graph.edge_count() {
            self.pheromone_graph[EdgeIndex::new(edge_no)] *= 1.0 - self.opt.evaporation;
            if new_path.contains(&edge_no) {
                self.pheromone_graph[EdgeIndex::new(edge_no)] +=
                    self.opt.density / (new_path.len() as f64);
            }
        }
        self.iteration += 1;

        if self.opt.verbose {
            self.export();
        }
    }

    /// Let the ants explore the problem
    pub fn explore(&mut self) -> (bool, f64, Vec<N>) {
        // Create a threadpool for this exploration
        let pool = ThreadPool::new(self.opt.cpus);

        let min_equal_paths = ((self.opt.ants as f64) * self.opt.threshold) as u32;

        while self.iteration < self.opt.iterations {
            let (tx, rx) = channel();

            for ant_no in 0..self.opt.ants {
                let tx = tx.clone();
                let colony = self.clone();
                let current_seed = self.init_seed;
                let objective = self.obj.clone();
                let start_idx = self.start_idx;
                pool.execute(move || {
                    // Let's sleep to _try_ to keep the ants ordered in the same way between runs
                    sleep(Duration::from_millis(5));

                    // Initialize the n-th Ant.
                    let mut ant = Ant::from_node(ant_no, start_idx, current_seed, objective);
                    while ant.move_ant(&colony) {}
                    ant.remove_cycles(&colony);
                    tx.send(ant)
                        .expect("channel will be there waiting for the pool");
                });

                self.init_seed += 1;
            }

            let mut shortest_path: Vec<usize> = Vec::new();
            let mut lowest_weight = 0.0;
            let mut equal_paths = 1;

            // And wait for all ants to finish
            rx.iter().take(self.opt.ants).for_each(|ant| {
                if ant.number == 0 || ant.edge_weight < lowest_weight {
                    lowest_weight = ant.edge_weight;
                    shortest_path = ant.visited_edges.clone();
                    equal_paths = 1; // Reset counter
                } else if (lowest_weight - ant.edge_weight).abs() < EPSILON {
                    // At least two ants have found the same path
                    equal_paths += 1;
                }
            });

            if equal_paths >= min_equal_paths && self.iteration > 0 {
                self.shortest_path = shortest_path;
                info!(
                    "Stagnation reached: {:0.2}% of ants have the same path",
                    self.opt.threshold * 100.0
                );
                return self.end(false);
            }

            self.update_pheromones(shortest_path);
            info!("Completed iteration {}", self.iteration);
        }
        self.end(true)
    }

    fn end(&self, max_iter: bool) -> (bool, f64, Vec<N>) {
        self.export();
        if max_iter {
            warn!("Max iterations reached")
        }
        // Compute the weight and path
        let mut weight = 0.0;
        let mut path = Vec::new();
        for edge_idx in self.shortest_path.clone() {
            let edge = EdgeIndex::new(edge_idx);
            let (start_nidx, end_nidx) = self.viz_graph.edge_endpoints(edge).unwrap();
            let start_node = self.viz_graph[start_nidx].clone();
            let end_node = self.viz_graph[end_nidx].clone();
            if path.is_empty() {
                if start_nidx == NodeIndex::new(self.start_idx) {
                    path.push(start_node);
                    path.push(end_node);
                } else {
                    // Add in reverse order
                    path.push(end_node);
                    path.push(start_node);
                }
            } else if path.last().unwrap() == &start_node {
                path.push(end_node);
            } else {
                path.push(start_node);
            }

            weight += self.viz_graph[edge];
        }

        println!(
            "iterations: {}\nweight: {}\npath: {:?}",
            self.iteration + 1,
            weight,
            path,
        );

        (!max_iter, weight, path)
    }

    fn export(&self) {
        export_dot(
            &self.pheromone_graph,
            &format!("dotgraphs/pheromones_it_{}.dot", self.iteration),
        );
    }
}

#[test]
fn test_colony_complete() {
    use std::path::PathBuf;
    use std::str::FromStr;
    let colony_opt = ColonyOpts {
        verbose: false,
        initphe: 1.0,
        evaporation: 0.95,
        iterations: 100,
        threshold: 1.0,
        density: 1.0,
        alpha: 1,
        beta: 2,
        seed: Some(4580147489518812583),
        visit_node: None,
        visit_count: None,
        reach_dead_end: false,
        start_node: None,
        wander: 0.15,
        ants: 10,
        cpus: 1,
    };

    let mut colony =
        Colony::with_input_path(colony_opt, PathBuf::from_str("complete_graph.txt").unwrap());

    let (success, weight, path) = colony.explore();

    assert!(success);
    assert!(weight - 8.1 < EPSILON);
    assert_eq!(path, vec!["S", "A", "D", "B", "C"]);
}