use crate::graph::Graph;
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use crate::structures::dsu::DSU;

impl <T> Graph<T> where T : PartialOrd + Copy + Default + std::ops::Add<Output = T> {
    pub fn kruskal(&self) -> (Graph<T>, T) {
        struct KruskalEdge<U> where U : PartialOrd + Copy {
            from: usize,
            to: usize,
            dist: U,
        }

        impl <U> std::cmp::PartialEq for KruskalEdge<U> where U : PartialOrd + Copy {
            fn eq(&self, other: &KruskalEdge<U>) -> bool {
                self.dist == other.dist
            }
        }

        impl <U> Eq for KruskalEdge<U> where U : PartialOrd + Copy {}

        impl <U> std::cmp::Ord for KruskalEdge<U> where U : PartialOrd + Copy {
            fn cmp(&self, other: &Self) -> Ordering {
                other.dist.partial_cmp(&self.dist).unwrap()
            }
        }

        impl <U> std::cmp::PartialOrd for KruskalEdge<U> where U : PartialOrd + Copy {
            fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
                Some(other.dist.partial_cmp(&self.dist).unwrap())
            }
        }

        let mut graph = Graph::new(self.adj.len());
        let mut heap = BinaryHeap::new();
        let mut dsu = DSU::new(graph.adj.len());
        for (from, edges) in self.adj.iter().enumerate().skip(1) {
            dsu.make_set(from).unwrap();
            for edge in edges {
                heap.push(KruskalEdge {
                    from,
                    to: edge.to,
                    dist: edge.weight
                });
            }
        }

        let mut weight = T::default();
        while let Some (value) = heap.pop() {
            if dsu.find_set(value.from) != dsu.find_set(value.to) {
                dsu.union_sets(value.from, value.to).unwrap();
                graph.add_edge(value.from, value.to, value.dist).unwrap();
                weight = weight + value.dist;
            }
        }
        (graph, weight)
    }
}

#[cfg(test)]
#[test]
fn test_kruskal(){
    let mut graph = Graph::new(20);

    graph.add_edge(1, 2, 7.0).unwrap();
    graph.add_edge(1, 4, 5.0).unwrap();
    graph.add_edge(2, 3, 8.0).unwrap();
    graph.add_edge(2, 5, 7.0).unwrap();
    graph.add_edge(2, 4, 9.0).unwrap();
    graph.add_edge(3, 5, 5.0).unwrap();
    graph.add_edge(5, 7, 9.0).unwrap();
    graph.add_edge(5, 6, 8.0).unwrap();
    graph.add_edge(5, 4, 15.0).unwrap();
    graph.add_edge(6, 7, 11.0).unwrap();
    graph.add_edge(6, 4, 6.0).unwrap();

    let (tree, weight) = graph.kruskal();
    assert_eq!(vec![1, 2, 5, 7], tree.search_path(7, &tree.bfs(1).unwrap()).unwrap().unwrap());
    assert_eq!(vec![1, 2, 5, 3], tree.search_path(3, &tree.bfs(1).unwrap()).unwrap().unwrap());
    assert_eq!(39.0, weight);
}