okvs 0.2.0

use std::collections::{hash_map::Entry, HashMap};

/// An edge with a value.
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Edge<V> {
    pub to_vertex: usize,
    pub value: V,
}

/// An undirected graph with a value on each edge.
#[derive(Debug, PartialEq)]
pub struct UndirectedGraph<V> {
    pub(crate) adjacency_list: Vec<Vec<Edge<V>>>,
    pub(crate) edge_count: usize,
    self_references: HashMap<usize, Edge<V>>,
}

impl<V: Copy> UndirectedGraph<V> {
    pub fn new(vertex_count: usize) -> Self {
        Self {
            adjacency_list: (0..vertex_count).map(|_| vec![]).collect(),
            edge_count: 0,
            self_references: HashMap::new(),
        }
    }

    pub fn has_edges(&self, vertex: usize) -> bool {
        if self.self_references.contains_key(&vertex) {
            return true;
        }

        !self.adjacency_list[vertex].is_empty()
    }

    /// Adds an edge between `vertex_a` and `vertex_b` with the specified `value`.
    /// Returns `true` if successful, and `false` if an edge already existed with a different value.
    pub fn add_edge(&mut self, vertex_a: usize, vertex_b: usize, value: V) -> bool {
        for edge in &self.adjacency_list[vertex_a] {
            if edge.to_vertex == vertex_b {
                return false;
            }
        }

        if vertex_a == vertex_b {
            if let Entry::Vacant(e) = self.self_references.entry(vertex_a) {
                e.insert(Edge {
                    to_vertex: vertex_b,
                    value,
                });
                self.edge_count += 1;
                return true;
            } else {
                return false;
            }
        }

        self.adjacency_list[vertex_a].push(Edge {
            to_vertex: vertex_b,
            value,
        });
        self.adjacency_list[vertex_b].push(Edge {
            to_vertex: vertex_a,
            value,
        });
        self.edge_count += 1;
        true
    }

    pub fn pop_edge(&mut self, vertex: usize) -> Option<Edge<V>> {
        if self.self_references.contains_key(&vertex) {
            self.edge_count -= 1;
            return self.self_references.remove(&vertex);
        }

        let edge = self.adjacency_list[vertex].pop()?;
        let index = self.adjacency_list[edge.to_vertex]
            .iter()
            .position(|edge| edge.to_vertex == vertex)
            .unwrap();
        self.adjacency_list[edge.to_vertex].swap_remove(index);
        self.edge_count -= 1;

        Some(edge)
    }

    /// Returns the edge if it is the only one remaining. A self-reference does not count, so the function returns `None`.
    pub fn pop_only_edge(&mut self, vertex: usize) -> Option<Edge<V>> {
        if self.adjacency_list[vertex].len() != 1 {
            return None;
        }

        if self.self_references.contains_key(&vertex) {
            return None;
        }

        let edge = self.adjacency_list[vertex].pop()?;
        let index = self.adjacency_list[edge.to_vertex]
            .iter()
            .position(|edge| edge.to_vertex == vertex)
            .unwrap();
        self.adjacency_list[edge.to_vertex].swap_remove(index);
        self.edge_count -= 1;

        Some(edge)
    }
}

#[cfg(test)]
mod tests {
    use super::UndirectedGraph;

    // TODO: Add testcases for self-references

    #[test]
    fn empty_graph() {
        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
        assert_eq!(graph.edge_count, 0);
        for i in 0..10 {
            assert!(graph.pop_only_edge(i).is_none());
        }
    }

    #[test]
    fn add_edge() {
        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
        graph.add_edge(5, 1, 2500);
        assert_eq!(graph.edge_count, 1);
    }

    #[test]
    fn add_edge_invariant() {
        let mut graph_a: UndirectedGraph<usize> = UndirectedGraph::new(10);
        graph_a.add_edge(5, 1, 3000);

        let mut graph_b: UndirectedGraph<usize> = UndirectedGraph::new(10);
        graph_b.add_edge(1, 5, 3000);

        assert_eq!(graph_a, graph_b);
    }

    #[test]
    fn pop_edge() {
        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
        graph.add_edge(5, 1, 2500);

        let edge_a = graph.pop_edge(1).unwrap();

        assert_eq!(graph.edge_count, 0);
        assert_eq!(edge_a.value, 2500);
        assert_eq!(edge_a.to_vertex, 5);
        assert!(graph.adjacency_list[1].is_empty());
        assert!(graph.adjacency_list[5].is_empty());

        graph.add_edge(5, 1, 2500);
        assert_eq!(graph.edge_count, 1);

        let edge_b = graph.pop_edge(5).unwrap();

        assert_eq!(graph.edge_count, 0);
        assert_eq!(edge_b.value, 2500);
        assert_eq!(edge_b.to_vertex, 1);
        assert!(graph.adjacency_list[1].is_empty());
        assert!(graph.adjacency_list[5].is_empty());
    }

    #[test]
    fn pop_only_edge() {
        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
        graph.add_edge(5, 1, 2500);

        let edge_a = graph.pop_only_edge(1).unwrap();

        assert_eq!(graph.edge_count, 0);
        assert_eq!(edge_a.value, 2500);
        assert_eq!(edge_a.to_vertex, 5);
        assert!(graph.adjacency_list[1].is_empty());
        assert!(graph.adjacency_list[5].is_empty());

        graph.add_edge(5, 1, 2500);
        assert_eq!(graph.edge_count, 1);

        let edge_b = graph.pop_only_edge(5).unwrap();

        assert_eq!(graph.edge_count, 0);
        assert_eq!(edge_b.value, 2500);
        assert_eq!(edge_b.to_vertex, 1);
        assert!(graph.adjacency_list[1].is_empty());
        assert!(graph.adjacency_list[5].is_empty());
    }
}