//! Region adjacency and the merge hierarchy — arm A's CPU half (S15).
//!
//! **Deliberately device-free.** Nothing here touches wgpu, and that is the
//! point: the hierarchy is where arm A's behaviour actually lives, so it has
//! to be assertable on hand-built inputs rather than only on whatever a GPU
//! happened to produce (ARCH §6.5a). Every test below runs on CI machines
//! with no adapter.
//!
//! It is also why this is a legitimate CPU stage rather than a violation of
//! ARCH §6.1. The watershed's basins are found on the GPU, over pixels; what
//! follows runs on the **region adjacency graph**, which for a 2 MP proxy is
//! a few thousand nodes. Pixels stay on the GPU, the graph is CPU-side —
//! the same split ARCH §3.4 already draws for the edit graph.
//!
//! # Why a hierarchy rather than one partition
//!
//! A single segmentation has exactly one granularity, and no threshold is
//! right for both "her eye" and "her face". So the watershed deliberately
//! *over*-segments, and the merge order recorded here is what lets one
//! interaction — a scroll, a drag — walk from a fragment up to the object
//! that contains it.

use std::collections::HashMap;

/// One boundary between two adjacent regions.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Edge {
    pub a: u32,
    pub b: u32,
    /// The **saddle**: the lowest point on the ridge separating the two
    /// basins, which is the height at which water joining them would first
    /// spill over.
    ///
    /// The minimum over the shared boundary, not the mean or the maximum. A
    /// long boundary that is strong everywhere except one weak gap describes
    /// two regions that should merge early — the gap is exactly where a
    /// person would say the edge fails.
    pub saddle: f32,
}

/// A partition of the image into labelled regions, plus how they adjoin.
///
/// The shared interface from docs/dev/segmentation.md §2: arm A produces this
/// from a watershed, arm B would produce it from a class map, and the
/// consumers above cannot tell which.
#[derive(Debug, Clone, PartialEq)]
pub struct RegionField {
    pub width: usize,
    pub height: usize,
    /// Region id per pixel, compacted to `0..region_count`.
    pub labels: Vec<u32>,
    pub region_count: usize,
    /// Deduplicated, `a < b`, sorted for reproducibility.
    pub adjacency: Vec<Edge>,
}

impl RegionField {
    /// Build from the watershed's raw output.
    ///
    /// `roots` holds, per pixel, the linear index of its basin root — what
    /// the pointer-jumping pass converges to. Those indices are sparse and
    /// arbitrary, so they are compacted here to `0..region_count` in order of
    /// first appearance, which makes them stable to serialise and cheap to
    /// index.
    pub fn from_roots(roots: &[u32], gradient: &[f32], width: usize, height: usize) -> Self {
        assert_eq!(roots.len(), width * height, "roots must cover every pixel");
        assert_eq!(
            gradient.len(),
            width * height,
            "gradient must cover every pixel"
        );

        // Compact in raster order rather than by hashing the root value, so
        // the same image always yields the same numbering (M5).
        let mut compact: HashMap<u32, u32> = HashMap::new();
        let mut labels = vec![0u32; roots.len()];
        for (i, &root) in roots.iter().enumerate() {
            let next = compact.len() as u32;
            labels[i] = *compact.entry(root).or_insert(next);
        }
        let region_count = compact.len();

        // Saddles, over 4-neighbour adjacency. Crossing a boundary means
        // climbing to the higher of the two pixels, so a pair's pass height
        // is the max; the region pair's saddle is the min over all its pairs.
        let mut saddles: HashMap<(u32, u32), f32> = HashMap::new();
        let note = |p: usize, q: usize, saddles: &mut HashMap<(u32, u32), f32>| {
            let (la, lb) = (labels[p], labels[q]);
            if la == lb {
                return;
            }
            let key = (la.min(lb), la.max(lb));
            let pass = gradient[p].max(gradient[q]);
            saddles
                .entry(key)
                .and_modify(|s| {
                    if pass < *s {
                        *s = pass;
                    }
                })
                .or_insert(pass);
        };

        for y in 0..height {
            for x in 0..width {
                let i = y * width + x;
                if x + 1 < width {
                    note(i, i + 1, &mut saddles);
                }
                if y + 1 < height {
                    note(i, i + width, &mut saddles);
                }
            }
        }

        let mut adjacency: Vec<Edge> = saddles
            .into_iter()
            .map(|((a, b), saddle)| Edge { a, b, saddle })
            .collect();
        // Sorted by (saddle, a, b) with `total_cmp` rather than `partial_cmp`:
        // a total order over floats, so the sequence cannot depend on how the
        // hash map happened to iterate. The merge order below is derived from
        // this, and an unstable merge order is an unstable label field.
        adjacency.sort_by(|x, y| {
            x.saddle
                .total_cmp(&y.saddle)
                .then(x.a.cmp(&y.a))
                .then(x.b.cmp(&y.b))
        });

        Self {
            width,
            height,
            labels,
            region_count,
            adjacency,
        }
    }

    /// Apply a cut's grouping to every pixel, for display or masking.
    pub fn apply(&self, grouping: &[u32]) -> Vec<u32> {
        self.labels.iter().map(|&l| grouping[l as usize]).collect()
    }
}

/// One merge in the hierarchy.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Merge {
    pub a: u32,
    pub b: u32,
    pub saddle: f32,
}

/// The merge order over a [`RegionField`] — the whole hierarchy.
///
/// Kruskal over the adjacency edges: sort by saddle, union-find, and record
/// each union that actually joined two distinct sets. The recording *is* the
/// dendrogram, which is why the multiscale part costs one `Vec` rather than a
/// second algorithm.
#[derive(Debug, Clone, PartialEq)]
pub struct MergeTree {
    /// In merge order, so saddles are non-decreasing.
    pub merges: Vec<Merge>,
    pub region_count: usize,
}

impl MergeTree {
    pub fn build(field: &RegionField) -> Self {
        let mut uf = UnionFind::new(field.region_count);
        let mut merges = Vec::new();

        // `field.adjacency` is already sorted by (saddle, a, b).
        for e in &field.adjacency {
            if uf.union(e.a as usize, e.b as usize) {
                merges.push(Merge {
                    a: e.a,
                    b: e.b,
                    saddle: e.saddle,
                });
            }
        }

        Self {
            merges,
            region_count: field.region_count,
        }
    }

    /// The grouping after applying every merge weaker than `threshold`.
    ///
    /// Returns one group id per original region, compacted to `0..groups`.
    pub fn cut_at(&self, threshold: f32) -> Vec<u32> {
        self.cut(self.merges.iter().take_while(|m| m.saddle <= threshold))
    }

    /// The grouping with at most `target` groups.
    ///
    /// The interface a scroll wheel wants: thresholds are in gradient units
    /// and mean nothing to anyone, where "about two hundred regions" is a
    /// granularity a person can ask for.
    pub fn cut_to(&self, target: usize) -> Vec<u32> {
        let target = target.max(1);
        let take = self.region_count.saturating_sub(target);
        self.cut(self.merges.iter().take(take))
    }

    fn cut<'a>(&self, merges: impl Iterator<Item = &'a Merge>) -> Vec<u32> {
        let mut uf = UnionFind::new(self.region_count);
        for m in merges {
            uf.union(m.a as usize, m.b as usize);
        }

        // Compact roots in region order, so a cut's numbering is stable and
        // independent of the merge sequence that produced it.
        let mut seen: HashMap<usize, u32> = HashMap::new();
        (0..self.region_count)
            .map(|r| {
                let root = uf.find(r);
                let next = seen.len() as u32;
                *seen.entry(root).or_insert(next)
            })
            .collect()
    }

    /// How many groups `cut_to` / `cut_at` would leave, without building one.
    pub fn groups_at(&self, threshold: f32) -> usize {
        let merged = self.merges.iter().filter(|m| m.saddle <= threshold).count();
        self.region_count - merged
    }
}

struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<u8>,
}

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

    fn find(&mut self, mut x: usize) -> usize {
        while self.parent[x] != x {
            // Path halving — the compression that keeps this near-linear
            // without the second pass full compression needs.
            self.parent[x] = self.parent[self.parent[x]];
            x = self.parent[x];
        }
        x
    }

    /// Returns whether this union joined two distinct sets.
    fn union(&mut self, a: usize, b: usize) -> bool {
        let (ra, rb) = (self.find(a), self.find(b));
        if ra == rb {
            return false;
        }
        // Union by rank, but with a deterministic tie-break: equal ranks
        // attach the higher index under the lower. Without it the tree shape
        // depends on argument order, and `find` would return a different
        // representative for the same input on a different day.
        let (lo, hi) = if self.rank[ra] > self.rank[rb] {
            (ra, rb)
        } else if self.rank[rb] > self.rank[ra] {
            (rb, ra)
        } else {
            let (lo, hi) = (ra.min(rb), ra.max(rb));
            self.rank[lo] += 1;
            (lo, hi)
        };
        self.parent[hi] = lo;
        true
    }
}

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

    /// A 4×2 field split down the middle, with a weak gap in the boundary.
    ///
    ///   labels    0 0 1 1      gradient   0 5 5 0
    ///             0 0 1 1                 0 2 2 0
    ///
    /// The saddle is 2, not 5: the bottom row is where the ridge is lowest.
    fn split_field() -> RegionField {
        let roots = vec![0, 0, 2, 2, 0, 0, 2, 2];
        let gradient = vec![0.0, 5.0, 5.0, 0.0, 0.0, 2.0, 2.0, 0.0];
        RegionField::from_roots(&roots, &gradient, 4, 2)
    }

    #[test]
    fn roots_compact_to_dense_ids() {
        // Basin roots are pixel indices and so are sparse and arbitrary.
        // Anything that indexes by region needs them dense.
        let f = split_field();
        assert_eq!(f.region_count, 2);
        assert_eq!(f.labels, vec![0, 0, 1, 1, 0, 0, 1, 1]);
    }

    #[test]
    fn the_saddle_is_the_lowest_pass_not_the_typical_one() {
        // The property that makes the hierarchy match human judgement: two
        // regions joined by one weak gap belong together, however strong the
        // rest of the boundary is.
        let f = split_field();
        assert_eq!(f.adjacency.len(), 1);
        assert_eq!(f.adjacency[0].saddle, 2.0);
    }

    #[test]
    fn a_uniform_image_is_one_region() {
        // No minima to separate, so nothing to merge — and the tree must not
        // invent a merge it cannot justify.
        let f = RegionField::from_roots(&[0; 16], &[0.0; 16], 4, 4);
        assert_eq!(f.region_count, 1);
        assert!(f.adjacency.is_empty());
        assert!(MergeTree::build(&f).merges.is_empty());
    }

    /// Three regions in a row: 0 | 1 | 2, with the 1–2 boundary weaker.
    fn chain_field() -> RegionField {
        let roots = vec![0, 0, 2, 2, 4, 4];
        let gradient = vec![0.0, 9.0, 0.0, 3.0, 0.0, 0.0];
        RegionField::from_roots(&roots, &gradient, 6, 1)
    }

    #[test]
    fn the_weakest_boundary_merges_first() {
        // The whole basis of the granularity ladder: walking up the tree must
        // dissolve the least convincing edge before a strong one.
        let tree = MergeTree::build(&chain_field());
        assert_eq!(tree.merges.len(), 2);
        assert_eq!(tree.merges[0].saddle, 3.0);
        assert_eq!(tree.merges[1].saddle, 9.0);
        assert_eq!((tree.merges[0].a, tree.merges[0].b), (1, 2));
    }

    #[test]
    fn merges_are_ordered_by_saddle() {
        let tree = MergeTree::build(&chain_field());
        assert!(
            tree.merges.windows(2).all(|w| w[0].saddle <= w[1].saddle),
            "a cut at a threshold is only meaningful if merges are ordered"
        );
    }

    #[test]
    fn a_cut_walks_from_every_region_to_one() {
        let tree = MergeTree::build(&chain_field());

        // The bottom of the ladder: nothing merged.
        let fine = tree.cut_to(3);
        assert_eq!(fine, vec![0, 1, 2]);

        // The middle: the weak boundary is gone, the strong one survives.
        let mid = tree.cut_to(2);
        assert_eq!(mid[1], mid[2], "the weak boundary should have dissolved");
        assert_ne!(mid[0], mid[1], "the strong boundary should survive");

        // The top: one region.
        let coarse = tree.cut_to(1);
        assert_eq!(coarse, vec![0, 0, 0]);
    }

    #[test]
    fn cut_to_is_clamped_rather_than_panicking() {
        // A scroll wheel runs past both ends of the ladder, and neither end
        // is an error.
        let tree = MergeTree::build(&chain_field());
        assert_eq!(tree.cut_to(0), tree.cut_to(1));
        assert_eq!(tree.cut_to(99), vec![0, 1, 2]);
    }

    #[test]
    fn cut_at_and_cut_to_agree() {
        let tree = MergeTree::build(&chain_field());
        // Above the weak saddle, below the strong one.
        assert_eq!(tree.cut_at(5.0), tree.cut_to(2));
        assert_eq!(tree.groups_at(5.0), 2);
    }

    #[test]
    fn a_grouping_maps_back_onto_pixels() {
        let f = chain_field();
        let tree = MergeTree::build(&f);
        let px = f.apply(&tree.cut_to(2));
        assert_eq!(px, vec![0, 0, 1, 1, 1, 1]);
    }

    #[test]
    fn the_same_input_gives_a_bit_identical_result() {
        // M5 in miniature. The CPU half must contribute no nondeterminism of
        // its own, or there is no point asking whether the GPU half does:
        // hash map iteration order is the obvious way to fail this, which is
        // why both the adjacency list and the cut numbering are sorted rather
        // than taken as the map yields them.
        let (a, b) = (chain_field(), chain_field());
        assert_eq!(a, b);

        let (ta, tb) = (MergeTree::build(&a), MergeTree::build(&b));
        assert_eq!(ta, tb);
        assert_eq!(ta.cut_to(2), tb.cut_to(2));
    }

    #[test]
    fn a_ladder_over_many_regions_is_monotone() {
        // Region counts must fall as you climb, with no level skipped —
        // otherwise a scroll step could jump past the granularity someone
        // wanted.
        let n = 32;
        let roots: Vec<u32> = (0..n).map(|i| i as u32).collect();
        let gradient: Vec<f32> = (0..n).map(|i| (i % 7) as f32).collect();
        let f = RegionField::from_roots(&roots, &gradient, n, 1);
        let tree = MergeTree::build(&f);

        let mut last = usize::MAX;
        for target in (1..=n).rev() {
            let groups = tree.cut_to(target).iter().max().map(|m| *m as usize + 1);
            let groups = groups.unwrap_or(0);
            assert_eq!(groups, target, "cut_to({target}) should leave {target}");
            assert!(groups < last);
            last = groups;
        }
    }
}