extern crate queues;

use super::Network;
use queues::*;
use rand::prelude::*;
use std::collections::HashSet;

/// Perform a breadth-first search over the network, starting with root and looking for target.
/// Return a HashSet of all nodes discovered in the process or `None` if the path is not found.
/// Returns `Err` if the root or target nodes do not exist. Panics if the network is corrupted or
/// when the queue operations fail.
/// # Examples
/// ```
/// let net = Network::new(100, NetworkModel::ER { p: 0.05, whole: false });
/// println!("{:?}\n", net);
/// match bfs(&net, 1, 3) {
///     Some(n) => println!("Found 3 after discovering {:?}", n),
///     None => println!("No path to 3 found!"),
/// }
/// ```
pub fn bfs(net: &Network, root: usize, target: usize) -> Result<Option<HashSet<usize>>, String> {
    // Check if the target is valid
    if let None = net.nodes.get(target) {
        return Err(format!("No such node in network: {}", target));
    }
    if let None = net.nodes.get(root) {
        return Err(format!("No such node in network: {}", root));
    }
    let mut discovered: HashSet<usize> = HashSet::new();
    let mut to_visit: Queue<usize> = Queue::new();
    discovered.insert(root);
    to_visit.add(root).expect("Queue error");
    // Search the network breadth-first
    while to_visit.size() != 0 {
        let current = to_visit.remove().expect("Queue error");
        // Exit condition
        if current == target {
            return Ok(Some(discovered));
        }
        for &i in net.nodes[current].links.keys() {
            if !discovered.contains(&i) {
                discovered.insert(i);
                to_visit.add(i).expect("Queue error");
            }
        }
    }
    // No path from root to target was found
    Ok(None)
}

/// Explore the network with BFS starting with `root`. Return a tuple of two HashSets: one with the
/// explored nodes (including the root) and another with the unexplored. Returns `Err` if the root
/// does not exist. Panics when the network is corrupted or if the queue operations fail.
/// # Examples
/// ```
/// let net = Network::new(100, NetworkModel::ER { p: 0.05, whole: false });
/// let (exp, unexp) = explore(&net, 1);
/// println!("Available from 1: {:?}", exp);
/// println!("Unavailable from 1: {:?}", unexp);
/// ```
pub fn explore(net: &Network, root: usize) -> Result<(HashSet<usize>, HashSet<usize>), String> {
    if let None = net.nodes.get(root) {
        return Err(format!("No such node: {}", root));
    }
    let mut discovered: HashSet<usize> = HashSet::new();
    let mut to_visit: Queue<usize> = Queue::new();
    discovered.insert(root);
    to_visit.add(root).expect("Queue error");
    // Explore the network just like BFS but without specified target
    while to_visit.size() != 0 {
        let current = to_visit.remove().expect("Queue error");
        for &i in net.nodes[current].links.keys() {
            if !discovered.contains(&i) {
                discovered.insert(i);
                to_visit.add(i).expect("Queue error");
            }
        }
    }
    // Gather the undiscovered nodes
    let mut undiscovered: HashSet<usize> = HashSet::new();
    for node in &net.nodes {
        if !discovered.contains(&node.index) {
            undiscovered.insert(node.index);
        }
    }
    Ok((discovered, undiscovered))
}

/// Stich the network together if it is not whole.
///
/// A random element of `net` is chosen as the root and the 'main' network is it's reachable
/// teritory. From the unreachable teritory a random node is chosen and connected to a random
/// elemnt of the 'main' network, thus marking it and its reachable teritory as no longer
/// unreachable. The process is repeated until there are no unreachable nodes left.
pub fn stitch_together(net: &mut Network) {
    let mut rng = rand::thread_rng();
    // Explore from a random node
    let root = net.nodes.choose(&mut rng).expect("Empty network").index;
    // Get connected and disconnected sets
    let (mut con, mut dc) = explore(net, root).expect("Failure when stitching");
    loop {
        // Choose a random disconnected node
        match dc.clone().into_iter().choose(&mut rng) {
            Some(nroot) => {
                // Establish its teritory (the 'new' teritory)
                let (n_teritory, _) = explore(net, nroot).expect("Failure when stitching");
                // Connect new root to a random element of con
                if let Some(r_con) = con.clone().into_iter().choose(&mut rng) {
                    // Con is never empty (it holds at least the root), hence ignore the `None`
                    net._link(nroot, r_con, &mut rng);
                }
                // Mark everything in the new teritory as connected
                for c in n_teritory {
                    con.insert(c);
                    dc.remove(&c);
                }
            }
            // The are no disconnected nodes
            None => break,
        }
    }
}


/// Performs the phase of information spreading according to the **susceptible-infected**
/// diffusion model.
///
/// On each time step the nodes for which the `infected` flag is `true` will try to infect
/// their closest neighbors with probability given by connection strength, setting their
/// `infected` flag to `true`.
pub fn spread_si(net: &mut Network, steps: usize) -> Result<usize, String> {
    // Choose random node as origin
    let mut rng = rand::thread_rng();
    let root = net.nodes.choose(&mut rng).expect("Empty network").index;
    net.nodes[root].infected = true;
    for _ in 0..steps {
        // Iterate over the copy of nodes as to not jumpstart infections
        for node in net.nodes.clone() {
            // In each time step try to infect connections
            // of already infected nodes with probability
            // given in the hashmap
            if node.infected {
                let map = node.links.clone();
                for &j in map.keys() {
                    if rng.gen::<f64>() < map[&j] {
                        net.nodes[j].infected = true;
                        }
                    }
                }
            }
        }
    Ok(root)
}