wybr 0.0.6

use std::collections::{HashMap, HashSet};

/// Outcome of extracting strongly connected components of a graph
pub struct StronglyConnected {
    index: u32,
    indices: HashMap<u32, u32>,
    lowlinks: HashMap<u32, u32>,
    on_stack: HashSet<u32>,
    stack: Vec<u32>,
    /// Each element is a strongly connected component, represented as a vector of graph vertex IDs
    pub components: Vec<Vec<u32>>,
    /// Edges of graph condensation
    pub cond_edges: HashSet<(usize, usize)>,
}

impl StronglyConnected {
    fn connect_one(&mut self, ev: u32, graph: &HashMap<u32, HashSet<u32>>) {
        use std::cmp::min;
        self.indices.insert(ev, self.index);
        self.lowlinks.insert(ev, self.index);
        self.index += 1;
        self.stack.push(ev);
        self.on_stack.insert(ev);
        if graph.contains_key(&ev) {
            for ew in graph.get(&ev).unwrap() {
                let n_ll = match (self.indices.get(ew).cloned(), self.on_stack.contains(ew)) {
                    (None, _) => {
                        self.connect_one(*ew, graph);
                        min(
                            *self.lowlinks.get(&ev).unwrap_or(&0),
                            *self.lowlinks.get(ew).unwrap_or(&0),
                        )
                    }
                    (Some(w_index), true) => min(*self.lowlinks.get(&ev).unwrap_or(&0), w_index),
                    _ => *self.lowlinks.get(&ev).unwrap_or(&0),
                };
                self.lowlinks.insert(ev, n_ll);
            }
        }

        if self.indices.get(&ev) == Some(self.lowlinks.get(&ev).unwrap_or(&0)) {
            let mut nc = Vec::new();
            while let Some(ee) = self.stack.pop() {
                self.on_stack.remove(&ee);
                nc.push(ee);
                if ee == ev {
                    break;
                }
            }
            //Not necessary, but allows checking output
            nc.sort_unstable();
            self.components.push(nc);
        }
    }
    fn connect(&mut self, graph: &HashMap<u32, HashSet<u32>>) {
        for e in graph.keys() {
            if !self.indices.contains_key(e) {
                self.connect_one(*e, graph)
            }
        }
    }
    fn condensate(&mut self, graph: &HashMap<u32, HashSet<u32>>) {
        let mut comp_map: HashMap<u32, usize> = HashMap::new();
        for (cid, el) in self.components.iter().enumerate() {
            for e in el {
                comp_map.insert(*e, cid);
            }
        }
        self.cond_edges.clear();
        for (f, dests) in graph {
            for &t in dests {
                let from = comp_map[f];
                let to = comp_map[&t];
                if from != to {
                    self.cond_edges.insert((from, to));
                }
            }
        }
    }
    /// Get a maximal element of the graph condensation
    ///
    /// Returns a vector of IDs of graph vertices which are in the union of all strongly connected
    /// components which are maximal elements of the graph condensation.
    pub fn get_max(&self) -> Vec<u32> {
        let mut max_i: HashSet<usize> = (0..self.components.len()).collect();
        for (_, t) in &self.cond_edges {
            max_i.remove(t);
        }
        max_i
            .iter()
            .flat_map(|&i| self.components[i].iter().cloned())
            .collect()
    }
}

/// Convert an iterator of edges into an outcome of strongly connected components search
///
/// Expects an iterator of `(f: u32, t: u32)`, which represent edge from vertex of an ID `f` to a vertex of
/// an id `t`.
/// Performs both strongly connected components identification with the Tarjan's algorithm and
/// graph condensation (to establish maximal elements required for Smith or Schwartz set.
use std::iter::FromIterator;
impl FromIterator<(u32, u32)> for StronglyConnected {
    fn from_iter<I: IntoIterator<Item = (u32, u32)>>(i: I) -> Self {
        let mut ans = StronglyConnected {
            index: 0,
            indices: HashMap::new(),
            lowlinks: HashMap::new(),
            on_stack: HashSet::new(),
            stack: Vec::new(),
            components: Vec::new(),
            cond_edges: HashSet::new(),
        };
        let mut graph: HashMap<u32, HashSet<u32>> = HashMap::new();
        for (from, to) in i {
            (*graph.entry(from).or_default()).insert(to);
        }
        ans.connect(&graph);
        ans.condensate(&graph);
        ans
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tally::VoteMatrix;

    #[test]
    fn sc_ex() {
        // 0->1, 1->2, 2->0, 3->4 3->2 4->0 --> 3 is a Condorcet Winner
        let ans: StronglyConnected = vec![(0, 1), (1, 2), (2, 0), (3, 4), (3, 2), (4, 0)]
            .iter()
            .map(|&(a, b)| (a, b))
            .collect();
        assert_eq!(ans.components, vec![vec![0, 1, 2], vec![4], vec![3]]);
    }
    #[test]
    fn sc_tricky() {
        let ans: StronglyConnected = vec![(0, 1), (3, 4), (1, 2), (4, 3), (2, 0), (1, 7), (4, 9)]
            .iter()
            .map(|&(a, b)| (a, b))
            .collect();
        let mut max = ans.get_max();
        max.sort_unstable();
        assert_eq!(max, vec![0, 1, 2, 3, 4]);
    }

    #[test]
    fn sc_vm() {
        // 0->1, 1->2, 2->0, 3->4 3->2 4->0 --> 3 is a Condorcet Winner
        let m = VoteMatrix::with_vector_bycol(&vec![
            0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1,
        ]);
        let ans: StronglyConnected = m
            .pairs()
            .filter_map(|((f, t), (w, o))| if w > o { Some((f, t)) } else { None })
            .collect();
        assert_eq!(ans.components, vec![vec![0, 1, 2], vec![4], vec![3]]);
    }
}