use crate::model::ImageId;
use std::collections::HashMap;

/// Disjoint-set union with path compression and union by rank.
pub struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<u8>,
}

impl UnionFind {
    pub fn new(n: usize) -> Self {
        UnionFind { parent: (0..n).collect(), rank: vec![0; n] }
    }

    pub fn find(&mut self, x: usize) -> usize {
        let mut root = x;
        while self.parent[root] != root {
            root = self.parent[root];
        }
        let mut cur = x;
        while self.parent[cur] != root {
            let next = self.parent[cur];
            self.parent[cur] = root;
            cur = next;
        }
        root
    }

    /// Merge the sets containing `a` and `b`. Returns true if they were distinct.
    pub fn union(&mut self, a: usize, b: usize) -> bool {
        let (ra, rb) = (self.find(a), self.find(b));
        if ra == rb {
            return false;
        }
        if self.rank[ra] < self.rank[rb] {
            self.parent[ra] = rb;
        } else if self.rank[ra] > self.rank[rb] {
            self.parent[rb] = ra;
        } else {
            self.parent[rb] = ra;
            self.rank[ra] += 1;
        }
        true
    }
}

/// A candidate match edge with its differing-pixel count as weight.
#[derive(Clone, Copy, Debug)]
pub struct CandidateEdge {
    pub a: ImageId,
    pub b: ImageId,
    pub diff: u64,
}

/// Compute a minimum spanning forest over `edges` (Kruskal), then orient each
/// tree's edges into parent pointers toward the smallest-id node (deterministic).
///
/// Returns a vector indexed by `ImageId`: `Some((parent_id, diff))` for a child,
/// `None` for a root. `edges` must already be filtered to the acceptance cap by
/// the caller.
pub fn spanning_forest(n: usize, mut edges: Vec<CandidateEdge>) -> Vec<Option<(ImageId, u64)>> {
    // Deterministic order: by weight, then by endpoint ids (normalized a<b).
    for e in edges.iter_mut() {
        if e.a > e.b {
            std::mem::swap(&mut e.a, &mut e.b);
        }
    }
    edges.sort_by_key(|e| (e.diff, e.a, e.b));

    // Kruskal: keep edges that join distinct components.
    let mut uf = UnionFind::new(n);
    let mut adj: HashMap<ImageId, Vec<(ImageId, u64)>> = HashMap::new();
    for e in &edges {
        if uf.union(e.a, e.b) {
            adj.entry(e.a).or_default().push((e.b, e.diff));
            adj.entry(e.b).or_default().push((e.a, e.diff));
        }
    }

    // Orient each tree from its smallest-id node via BFS.
    let mut parents: Vec<Option<(ImageId, u64)>> = vec![None; n];
    let mut visited = vec![false; n];
    for root in 0..n {
        if visited[root] {
            continue;
        }
        visited[root] = true;
        let mut queue = std::collections::VecDeque::new();
        queue.push_back(root);
        while let Some(node) = queue.pop_front() {
            if let Some(neighbours) = adj.get(&node) {
                // Sort for deterministic traversal.
                let mut ns = neighbours.clone();
                ns.sort_by_key(|(nbr, _)| *nbr);
                for (nbr, diff) in ns {
                    if !visited[nbr] {
                        visited[nbr] = true;
                        parents[nbr] = Some((node, diff));
                        queue.push_back(nbr);
                    }
                }
            }
        }
    }
    parents
}

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

    fn edge(a: usize, b: usize, diff: u64) -> CandidateEdge {
        CandidateEdge { a, b, diff }
    }

    #[test]
    fn union_find_merges_and_reports() {
        let mut uf = UnionFind::new(3);
        assert!(uf.union(0, 1));
        assert!(!uf.union(0, 1)); // already same set
        assert_eq!(uf.find(0), uf.find(1));
        assert_ne!(uf.find(0), uf.find(2));
    }

    #[test]
    fn chain_orients_from_smallest_root() {
        // A-B diff 5, B-C diff 1. MST keeps both. Root = id 0.
        let parents = spanning_forest(3, vec![edge(0, 1, 5), edge(1, 2, 1)]);
        assert_eq!(parents[0], None); // root
        assert_eq!(parents[1], Some((0, 5)));
        assert_eq!(parents[2], Some((1, 1)));
    }

    #[test]
    fn separate_components_each_have_a_root() {
        // Only B-C present; A (id 0) is isolated.
        let parents = spanning_forest(3, vec![edge(1, 2, 1)]);
        assert_eq!(parents[0], None);
        assert_eq!(parents[1], None); // smallest in its component
        assert_eq!(parents[2], Some((1, 1)));
    }

    #[test]
    fn parent_can_have_multiple_children() {
        // Triangle, all diff 1. MST keeps two edges; B-C is redundant.
        let parents = spanning_forest(3, vec![edge(0, 1, 1), edge(0, 2, 1), edge(1, 2, 1)]);
        assert_eq!(parents[0], None);
        assert_eq!(parents[1], Some((0, 1)));
        assert_eq!(parents[2], Some((0, 1))); // both children point at 0
    }

    #[test]
    fn cheapest_edges_chosen_first() {
        // A-B diff 9, A-C diff 1, B-C diff 1. MST should pick the two diff-1 edges.
        let parents = spanning_forest(3, vec![edge(0, 1, 9), edge(0, 2, 1), edge(1, 2, 1)]);
        // Root 0; 0-2 (diff1) and then 1-2 (diff1) connect everyone; 0-1(9) skipped.
        assert_eq!(parents[0], None);
        assert_eq!(parents[2], Some((0, 1)));
        assert_eq!(parents[1], Some((2, 1)));
    }

    #[test]
    fn output_is_independent_of_input_edge_order() {
        let forward = vec![edge(0, 2, 1), edge(1, 2, 1), edge(0, 1, 9)];
        let reversed: Vec<_> = forward.iter().rev().copied().collect();
        let shuffled = vec![edge(1, 2, 1), edge(0, 1, 9), edge(0, 2, 1)];
        let a = spanning_forest(3, forward);
        let b = spanning_forest(3, reversed);
        let c = spanning_forest(3, shuffled);
        assert_eq!(a, b);
        assert_eq!(a, c);
    }

    #[test]
    fn normalizes_reversed_endpoints() {
        // An edge given as (2, 1) must behave identically to (1, 2).
        let parents = spanning_forest(3, vec![edge(2, 1, 1)]);
        assert_eq!(parents[0], None);
        assert_eq!(parents[1], None); // smallest in its component is the root
        assert_eq!(parents[2], Some((1, 1)));
    }

    #[test]
    fn roots_at_smallest_id_when_component_excludes_zero() {
        // Component {3,4,5} (0,1,2 isolated) must root at 3.
        let parents = spanning_forest(6, vec![edge(4, 5, 1), edge(3, 4, 1)]);
        assert_eq!(parents[3], None);
        assert_eq!(parents[4], Some((3, 1)));
        assert_eq!(parents[5], Some((4, 1)));
    }
}
