DarkRoom/core/dr-segment/src/refine.rs

//! Sharpening a coarse category mask against the photograph's own colours.
//!
//! # The problem, as a number
//!
//! The `scene` module keeps the model's native `[1, 150, 80, 80]` logit grid,
//! so one cell is eight of the graph's input pixels across. At the 1600px
//! proxy the application segments at, the letterbox scale is `0.4` and **one
//! cell is 20 proxy pixels** — which the bilinear in `Scene::rasterise` then
//! spreads across one more either side.
//!
//! So a flag in the sky, or a chimney, or a bare branch, sits inside a handful
//! of cells whose softmax is dominated by the sky around it, and comes out
//! weighted as sky. No feather setting recovers it, because the information
//! was never in the grid.
//!
//! # What this does about it
//!
//! The photograph is at full proxy resolution even though the weights are not,
//! and it knows exactly where the flag is. The move is to let the *model*
//! decide what the category is and the *pixels* decide which of them belong to
//! it — the same division of labour [`crate::prior`] already draws between the
//! instance model and the watershed, applied to categories.
//!
//! 1. **Seeds.** Threshold the weights high, then erode by
//!    [`RefineOptions::margin_cells`] logit cells using
//!    [`crate::distance::signed_distance`]. What survives is confidently
//!    inside; the mirror of it is confidently outside. The erosion radius is
//!    derived from the grid rather than picked, because one cell *is* the
//!    model's resolution and anything within a cell of the boundary is
//!    precisely what there is no evidence about.
//!
//! 2. **A distribution per side, with several modes.** A single Gaussian over
//!    sky is wrong: sky is blue at the zenith, white where the cloud is, and
//!    pale at the horizon, and one blob over all three rejects two of them. So
//!    k-means, which is GrabCut's mixture without the EM.
//!
//! 3. **A mode is charged for its rarity.** A *small* flag deep in the sky has
//!    both a high weight and a large distance from the mask's boundary, so it
//!    lands in the interior sample and fits itself a mode. It cannot be
//!    excluded geometrically. What separates it from a real appearance of the
//!    category is that it explains almost none of the category — so a colour
//!    whose only explanation is a thin mode is charged `−ln(share)` nats for
//!    it, and a flag ends up far less plausible than the sky around it without
//!    ever being ruled out by fiat.
//!
//! 4. **Two tests, and a pixel must pass both.** An absolute one — is this
//!    colour plausible under the category at all, as a chi-square on the
//!    Mahalanobis distance — and a comparative one, is it likelier under the
//!    interior than the exterior. The absolute test is what catches the flag,
//!    whose colour is far from *both* sides and which the comparative test
//!    alone would leave at even odds. The comparative test is what stops the
//!    absolute one from needing a per-category constant.
//!
//! # Why the verdict is kept, and the decision is not
//!
//! Everything above is per-image and costs a distance transform, a k-means and
//! a pass over the pixels. What comes out of it is one number per pixel — how
//! much better the category explains this colour than the alternatives — and
//! [`Refinement`] stores exactly that, quantised.
//!
//! The *decision* is then a smoothstep on that number, which is arithmetic. So
//! [`Refinement::apply`] is cheap enough to run under a dragging slider, and
//! the model is never consulted again.
//!
//! This is [`crate::distance`]'s arrangement, deliberately: there, a signed
//! distance field is computed once and feather, grow and shrink become
//! arithmetic on it, "which is what makes those live controls rather than ones
//! that stall on every drag". Same shape, different field.
//!
//! # What the strictness control means
//!
//! [`Refinement::apply`] takes a strictness in nats, and it is a real quantity
//! rather than an arbitrary 0–100: it is how much more plausible than the
//! alternatives a colour must be before its weight is kept.
//!
//! Because a mode's rarity is *added* to the evidence against it, the things a
//! photographer wants to remove come off in the order they want them to. A
//! flag holding two percent of the sky is charged several nats; a bank of
//! cloud holding a third of it is charged none. So the strictness that takes
//! out the flag leaves the cloud, and the cloud only starts to go some way
//! further up — which is the property that makes a single slider worth
//! offering. It is pinned by `the_flag_goes_before_the_cloud_does`.
//!
//! At zero the refinement does nothing at all, which is what makes it safe as
//! a default-on feature: the control has an off position that is exactly the
//! old behaviour.
//!
//! # Two properties that make it safe to apply blind
//!
//! **It is subtractive.** The output is the input times a factor in `0..=1`,
//! never more. So the worst failure available to it is losing part of a real
//! sky; it can never *gain* a region, and a blue car below the horizon that
//! was never in the mask cannot be pulled into it.
//!
//! **The partition survives.** `scene.rs` rests entirely on the categories
//! summing to at most one at every pixel — that is what lets two adjacent
//! grades be feathered without painting both into the overlap. Multiplying by
//! a factor in `0..=1` cannot raise a sum, so refining every category
//! independently still leaves a partition. The weight taken off the flag lands
//! in the unlisted remainder, which is exactly where a flag belongs: ADE20K
//! has no class for one.
//!
//! # And when it cannot tell
//!
//! Every path that lacks the evidence to judge refuses to build a
//! [`Refinement`] at all and says why, rather than returning a plausible one.
//! An empty seed set fitted to a distribution would reject *every* pixel, and
//! a sky that silently vanished from a mask is the kind of failure nobody
//! attributes to the right place.

use crate::distance::signed_distance;

/// How the evidence for a refinement is gathered.
///
/// Everything here feeds the *fit*, which happens once per image. The one
/// control that moves afterwards is the strictness passed to
/// [`Refinement::apply`], and it is deliberately not in this struct: mixing
/// the two would invite a caller to change a fit parameter under a slider and
/// wonder why the drag stalled.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RefineOptions {
    /// Weight at or above which a pixel may seed the interior distribution.
    ///
    /// High on purpose. This is not the mask's threshold — it is the standard
    /// for "the model was sure here", and everything downstream is a
    /// consequence of what these pixels look like.
    pub seed_weight: f32,

    /// How far inside the boundary a seed must lie, in logit cells.
    ///
    /// One cell is the model's own resolution and `rasterise`'s bilinear
    /// spreads a cell's influence across one more, so the first cell and a
    /// half either side of the boundary is smear rather than evidence.
    pub margin_cells: f32,

    /// Modes fitted per side.
    ///
    /// Rarity is measured against an even split ([`Mode::rarity`]), so this
    /// number does not shift the verdicts on its own — raising it buys a finer
    /// description of the category without making every colour look rarer.
    pub clusters: usize,

    /// How much luminance counts against chrominance in judging a colour.
    ///
    /// Well under one, and that is the difference between this working and not
    /// on sky. Sky's variance is *dominated* by luminance — the zenith is
    /// several stops off the horizon and a cloud is brighter than either — so
    /// at equal weight the distribution is a long bright streak that a
    /// mid-grey flag sits comfortably inside. A flag is separated by
    /// chrominance; a cloud is separated by luminance alone. Down-weighting
    /// keeps the cloud and catches the flag.
    ///
    /// Not zero, or a dark bird against a bright sky is kept.
    ///
    /// The same reasoning `dr_gpu::SegmentOptions` applies to the watershed's
    /// gradient, for the same reason.
    pub luma_weight: f32,

    /// Squared Mahalanobis distance at which a colour stops being plausible
    /// under the category.
    ///
    /// Chi-square with three degrees of freedom: `11.34` is the 99th
    /// percentile, so under the fitted model one confident pixel in a hundred
    /// is expected to fail this on its own. It sets where the verdict scale's
    /// zero falls; the strictness then moves the decision along that scale.
    pub outlier: f32,

    /// Width of the band, in nats, over which the gate falls from keep to cut.
    ///
    /// A hard cut would put an aliased edge through the photograph at exactly
    /// the place a person is looking.
    pub soften: f32,

    /// Radius, in pixels, the *verdict* is blurred by.
    ///
    /// A per-pixel colour test on a real sky speckles: sensor noise puts
    /// individual pixels over the line in both directions. Blurring the
    /// evidence rather than the decision means it is paid for once, in
    /// [`Refinement::compute`], instead of on every frame of a drag.
    pub smooth_px: f32,

    /// Samples below which a side is too small to fit anything to.
    pub min_samples: usize,
}

impl Default for RefineOptions {
    fn default() -> Self {
        Self {
            seed_weight: 0.9,
            margin_cells: 1.5,
            clusters: 4,
            luma_weight: 0.25,
            outlier: 11.34,
            soften: 1.5,
            smooth_px: 2.0,
            min_samples: 500,
        }
    }
}

/// Strictness at which a category is left exactly as the model weighted it.
pub const STRICTNESS_OFF: f32 = 0.0;

/// The top of the useful strictness range, in nats.
///
/// Past this a category's own dominant colours have gone, so a control
/// offering more would only offer ways to delete the mask. Fixed here rather
/// than in the UI so that the slider and the tests agree on what the end of
/// the scale means.
pub const STRICTNESS_MAX: f32 = 8.0;

/// Where the control sits until someone moves it.
///
/// Half scale, and that position is the measurement rather than a round
/// number. On the synthetic frame `the_flag_goes_before_the_cloud_does`
/// builds, a flag holding 1.6% of the sky is more than half gone by **2.95**
/// and a cloud bank holding a third of it survives to **5.75** — so the whole
/// of the interval between them is a place the control can sit, and this is
/// near its middle.
///
/// That separation is what rarity buys. Without it both land at the same
/// verdict, there is no interval, and the control would appear dead until it
/// suddenly ate the sky.
///
/// Treat the numbers as a calibration and not a law: the synthetic sky is
/// smooth enough to sit on [`VARIANCE_FLOOR`], where a real one has noise and
/// therefore a real spread, which moves every crossing down together. The
/// ordering survives that; the exact placement is what the slider is for.
pub const STRICTNESS_DEFAULT: f32 = 4.0;

/// Why a refinement could not be built.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SkipReason {
    /// Nothing survived the erosion — the category is present but everywhere
    /// thinner than the model's own resolution, so there is no pixel it is
    /// sure about.
    NoInterior,
    /// The frame is essentially all this category, leaving nothing to
    /// contrast it against.
    NoExterior,
    /// The buffers handed in do not describe one image.
    Mismatched,
}

/// What a one-shot [`refine_category`] did.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Refined {
    /// Applied. `removed` is the fraction of the category's original total
    /// weight the gate took away.
    Applied {
        removed: f32,
    },
    Skipped(SkipReason),
}

/// Number of features a pixel is judged on: one luminance, two chrominance.
const FEATURES: usize = 3;

/// Floor on a fitted variance.
///
/// A mode over a flat patch of sky has a variance of nearly nothing, and
/// nothing in the denominator makes every other pixel infinitely improbable.
/// This is a standard deviation of one part in a hundred, comfortably below
/// anything a photograph resolves and comfortably above zero.
const VARIANCE_FLOOR: f32 = 1e-4;

/// Cap on the samples fitted per side.
///
/// k-means over a million pixels answers the same as k-means over twenty
/// thousand of them, and costs fifty times as much. The subsample is a fixed
/// stride rather than a random draw so the result is reproducible.
const MAX_SAMPLES: usize = 20_000;

/// Lloyd iterations.
///
/// Fixed rather than run to convergence: a stopping rule that depends on
/// floating-point comparison is a stopping rule that can differ between
/// machines, and a mask that differs between machines reaches the sidecar as
/// indices meaning one thing on the desktop and another on the phone
/// (docs/segmentation.md §6).
const ITERATIONS: usize = 12;

/// Half-width of the verdict scale, in nats.
///
/// The verdict is stored as a byte, so it needs an end. Sixteen is comfortably
/// past [`STRICTNESS_MAX`] plus a soften band, and it puts the quantisation
/// step at an eighth of a nat — an order finer than the narrowest transition
/// the gate can be asked for.
const VERDICT_RANGE: f32 = 16.0;

/// The evidence for refining one category, ready for a control to act on.
///
/// Built once per image by [`Refinement::compute`]; read by
/// [`Refinement::apply`] as often as a slider moves.
#[derive(Debug, Clone, PartialEq)]
pub struct Refinement {
    /// Per pixel, how much better the category explains this colour than the
    /// alternatives, in nats, quantised over ±[`VERDICT_RANGE`].
    ///
    /// A byte for the same reason a coverage buffer is one: at four bytes a
    /// pixel this would be a proxy-sized `f32` buffer per category, and the
    /// quantisation is far finer than the decision it feeds.
    verdict: Vec<u8>,
    width: usize,
    height: usize,
    /// Carried from the fit so that a caller holding only this can apply it
    /// without also having to keep the options that produced it.
    soften: f32,
}

impl Refinement {
    /// Fit the colour models and score every pixel. The expensive half.
    ///
    /// `weights` is one category's coverage at `width * height`, as
    /// `Scene::rasterise` produces it. `rgb` is the same picture, tightly
    /// packed `f32` RGB — the proxy the model itself read, so the two describe
    /// one frame.
    ///
    /// `cell_pixels` is how many pixels of *this* buffer one logit cell spans;
    /// see `Scene::cell_pixels`, which is where the caller should get it
    /// rather than re-deriving a letterbox inverse.
    pub fn compute(
        weights: &[f32],
        rgb: &[f32],
        width: usize,
        height: usize,
        cell_pixels: f32,
        options: &RefineOptions,
    ) -> Result<Self, SkipReason> {
        let pixels = width * height;
        if weights.len() != pixels || rgb.len() != pixels * 3 || pixels == 0 {
            return Err(SkipReason::Mismatched);
        }

        // The distance field is over the *thresholded* weights, so the
        // boundary it measures from is where the model stopped being sure —
        // not where the mask will eventually be cut, which is what the
        // strictness decides later.
        let coverage: Vec<u8> = weights
            .iter()
            .map(|&w| (w.clamp(0.0, 1.0) * 255.0).round() as u8)
            .collect();
        let threshold = (options.seed_weight.clamp(0.0, 1.0) * 255.0).round() as u8;
        let distance = signed_distance(&coverage, width, height, threshold);

        let margin = options.margin_cells * cell_pixels.max(1.0);

        let interior = sample(&distance, rgb, options.luma_weight, |d| d >= margin);
        if interior.len() < options.min_samples {
            return Err(SkipReason::NoInterior);
        }
        let exterior = sample(&distance, rgb, options.luma_weight, |d| d <= -margin);
        if exterior.len() < options.min_samples {
            return Err(SkipReason::NoExterior);
        }

        let inside = fit(&interior, options.clusters);
        let outside = fit(&exterior, options.clusters);

        let mut verdict = vec![0.0f32; pixels];
        for (p, cell) in verdict.iter_mut().enumerate() {
            let x = feature(rgb, p, options.luma_weight);

            // Squared Mahalanobis for the absolute test, because that is what
            // is chi-square distributed; the density — the same distance with
            // the mode's own spread folded in — for the comparative one, so
            // that a tight mode and a loose one are compared fairly. Both
            // carry the mode's rarity.
            let (maha_in, nll_in) = nearest(&inside, &x);
            let (_, nll_out) = nearest(&outside, &x);

            let plausible = 0.5 * (options.outlier - maha_in);
            let likelier = nll_out - nll_in;
            *cell = plausible.min(likelier);
        }

        // The evidence is blurred, not the decision — see
        // `RefineOptions::smooth_px`.
        if options.smooth_px >= 1.0 {
            verdict = blur(&verdict, width, height, options.smooth_px.round() as usize);
        }

        Ok(Self {
            verdict: verdict.iter().map(|&v| quantise(v)).collect(),
            width,
            height,
            soften: options.soften,
        })
    }

    pub fn size(&self) -> (usize, usize) {
        (self.width, self.height)
    }

    /// The gate this refinement applies at `strictness`, per pixel.
    ///
    /// Separate from [`Self::apply`] because the overlay wants to draw it and
    /// the coverage path wants to multiply by it, and neither should have to
    /// reimplement the smoothstep.
    pub fn gate(&self, strictness: f32) -> Vec<f32> {
        let strictness = strictness.max(0.0);
        self.verdict
            .iter()
            .map(|&v| smoothstep(-self.soften, self.soften, dequantise(v) - strictness))
            .collect()
    }

    /// Apply at a given strictness. The cheap half — a slider drags on this.
    ///
    /// At [`STRICTNESS_OFF`] this is the identity, exactly, and returns the
    /// weights unchanged without touching the verdict at all. That is what
    /// makes the control's zero position mean "as the model weighted it"
    /// rather than "very nearly that".
    pub fn apply(&self, weights: &[f32], strictness: f32) -> Vec<f32> {
        if strictness <= STRICTNESS_OFF || weights.len() != self.verdict.len() {
            return weights.to_vec();
        }
        weights
            .iter()
            .zip(self.gate(strictness))
            .map(|(&w, g)| w * g)
            .collect()
    }

    /// [`Self::apply`] over a quantised coverage buffer.
    ///
    /// The form the UI actually holds a category in: the mask reaches the
    /// rasteriser as one byte per pixel, and round-tripping it through floats
    /// to reuse `apply` would allocate twice as much for no extra precision.
    pub fn apply_coverage(&self, coverage: &[u8], strictness: f32) -> Vec<u8> {
        if strictness <= STRICTNESS_OFF || coverage.len() != self.verdict.len() {
            return coverage.to_vec();
        }
        coverage
            .iter()
            .zip(self.gate(strictness))
            .map(|(&c, g)| (c as f32 * g).round().clamp(0.0, 255.0) as u8)
            .collect()
    }
}

/// TRACES: FR-DEV-3
/// Fit and apply in one call, for a caller with no control to offer.
///
/// The example and the tests use this. The application does not: it keeps the
/// [`Refinement`] so that its slider costs a smoothstep rather than a k-means.
pub fn refine_category(
    weights: &[f32],
    rgb: &[f32],
    width: usize,
    height: usize,
    cell_pixels: f32,
    strictness: f32,
    options: &RefineOptions,
) -> (Vec<f32>, Refined) {
    match Refinement::compute(weights, rgb, width, height, cell_pixels, options) {
        Ok(refinement) => {
            let refined = refinement.apply(weights, strictness);
            let before: f32 = weights.iter().sum();
            let after: f32 = refined.iter().sum();
            let removed = if before > 0.0 {
                ((before - after) / before).clamp(0.0, 1.0)
            } else {
                0.0
            };
            (refined, Refined::Applied { removed })
        }
        Err(why) => (weights.to_vec(), Refined::Skipped(why)),
    }
}

fn quantise(v: f32) -> u8 {
    let t = (v.clamp(-VERDICT_RANGE, VERDICT_RANGE) + VERDICT_RANGE) / (2.0 * VERDICT_RANGE);
    (t * 255.0).round() as u8
}

fn dequantise(b: u8) -> f32 {
    (b as f32 / 255.0) * 2.0 * VERDICT_RANGE - VERDICT_RANGE
}

/// One pixel's colour, as the three numbers the distributions are fitted over.
///
/// Opponent axes rather than raw RGB, and no division anywhere: normalising
/// chrominance by luminance would be exposure-invariant and would also blow up
/// in the shadows, which on a landscape is the whole foreground.
fn feature(rgb: &[f32], p: usize, luma_weight: f32) -> [f32; FEATURES] {
    let (r, g, b) = (rgb[p * 3], rgb[p * 3 + 1], rgb[p * 3 + 2]);
    let luma = 0.2126 * r + 0.7152 * g + 0.0722 * b;
    [luma_weight * luma, r - g, b - 0.5 * (r + g)]
}

/// Every pixel whose distance passes `want`, subsampled to [`MAX_SAMPLES`].
///
/// Two passes — count, then take a fixed stride — rather than one pass
/// collecting everything and thinning afterwards, so a sky filling a 22 MP
/// proxy never allocates twenty million features to throw most of them away.
fn sample(
    distance: &[f32],
    rgb: &[f32],
    luma_weight: f32,
    want: impl Fn(f32) -> bool,
) -> Vec<[f32; FEATURES]> {
    let total = distance.iter().filter(|&&d| want(d)).count();
    if total == 0 {
        return Vec::new();
    }
    let stride = total.div_ceil(MAX_SAMPLES).max(1);

    let mut out = Vec::with_capacity(total.div_ceil(stride));
    let mut seen = 0usize;
    for (p, &d) in distance.iter().enumerate() {
        if !want(d) {
            continue;
        }
        if seen.is_multiple_of(stride) {
            out.push(feature(rgb, p, luma_weight));
        }
        seen += 1;
    }
    out
}

/// One mode of a side's colour distribution.
#[derive(Debug, Clone, Copy)]
struct Mode {
    mean: [f32; FEATURES],
    /// Per-feature, not a full covariance. The opponent axes are close enough
    /// to decorrelated by construction that the off-diagonal terms buy little,
    /// and a diagonal fit stays well-conditioned on the few hundred samples a
    /// small mode gets.
    variance: [f32; FEATURES],
    /// `0.5 * ln|Σ|`, precomputed because it is constant per mode and the
    /// inner loop runs once per pixel.
    half_log_det: f32,
    /// What it costs, in nats, to explain a colour only by this mode.
    ///
    /// `−ln(share × k)`, floored at zero. Three things follow from that shape:
    ///
    /// - It is measured against an **even split**, so raising
    ///   [`RefineOptions::clusters`] describes the category more finely
    ///   without making every colour in it look rarer.
    /// - It is floored, so a dominant mode earns no *bonus*. A discount there
    ///   would let the category's commonest colour outvote a genuinely bad
    ///   chi-square, which is the one direction this must not bend.
    /// - It is what replaces pruning small modes outright. A flag holding two
    ///   percent of the sky is charged a few nats rather than deleted, so
    ///   where it falls relative to the cloud is a *distance* the strictness
    ///   control can travel across, instead of a threshold it jumps over.
    rarity: f32,
}

/// Fit `k` modes to one side.
fn fit(samples: &[[f32; FEATURES]], k: usize) -> Vec<Mode> {
    let k = k.max(1).min(samples.len());
    let mut centres = seed_centres(samples, k);
    let k = centres.len();

    let mut assignment = vec![0usize; samples.len()];
    for _ in 0..ITERATIONS {
        let mut moved = false;
        for (s, sample) in samples.iter().enumerate() {
            let mut best = 0usize;
            let mut best_d = f32::INFINITY;
            for (c, centre) in centres.iter().enumerate() {
                let d = squared(sample, centre);
                if d < best_d {
                    best_d = d;
                    best = c;
                }
            }
            if assignment[s] != best {
                assignment[s] = best;
                moved = true;
            }
        }
        if !moved {
            break;
        }

        let mut sums = vec![[0.0f32; FEATURES]; k];
        let mut counts = vec![0usize; k];
        for (s, sample) in samples.iter().enumerate() {
            let c = assignment[s];
            for f in 0..FEATURES {
                sums[c][f] += sample[f];
            }
            counts[c] += 1;
        }
        for c in 0..k {
            // An emptied centre is left where it was rather than re-seeded.
            // Re-seeding is the usual advice, but it needs a random draw
            // inside the loop and the determinism is worth more here than the
            // empty cluster costs.
            if counts[c] == 0 {
                continue;
            }
            for f in 0..FEATURES {
                centres[c][f] = sums[c][f] / counts[c] as f32;
            }
        }
    }

    // Mean, spread and share from the final assignment.
    let mut sums = vec![[0.0f32; FEATURES]; k];
    let mut squares = vec![[0.0f32; FEATURES]; k];
    let mut counts = vec![0usize; k];
    for (s, sample) in samples.iter().enumerate() {
        let c = assignment[s];
        for f in 0..FEATURES {
            sums[c][f] += sample[f];
            squares[c][f] += sample[f] * sample[f];
        }
        counts[c] += 1;
    }

    let total = samples.len() as f32;
    let even = k as f32;
    let mut modes = Vec::new();
    for c in 0..k {
        if counts[c] == 0 {
            continue;
        }
        let n = counts[c] as f32;
        let mut mean = [0.0f32; FEATURES];
        let mut variance = [0.0f32; FEATURES];
        let mut half_log_det = 0.0f32;
        for f in 0..FEATURES {
            mean[f] = sums[c][f] / n;
            variance[f] = (squares[c][f] / n - mean[f] * mean[f]).max(VARIANCE_FLOOR);
            half_log_det += 0.5 * variance[f].ln();
        }
        modes.push(Mode {
            mean,
            variance,
            half_log_det,
            rarity: -((n / total) * even).ln().min(0.0),
        });
    }
    modes
}

/// k-means++ seeding, with a fixed sequence.
///
/// The usual seeding is random, and random would make a mask depend on which
/// run produced it. The draw is a fixed-seed LCG instead: still spread out,
/// still far better than taking the first `k` samples, and the same every time
/// on every machine.
fn seed_centres(samples: &[[f32; FEATURES]], k: usize) -> Vec<[f32; FEATURES]> {
    let mut rng = Lcg(0x2545_F491_4F6C_DD1D);
    let mut centres = vec![samples[0]];
    let mut nearest = vec![f32::INFINITY; samples.len()];

    while centres.len() < k {
        let last = *centres.last().expect("seeded with one");
        let mut total = 0.0f64;
        for (s, sample) in samples.iter().enumerate() {
            nearest[s] = nearest[s].min(squared(sample, &last));
            total += nearest[s] as f64;
        }
        if total <= 0.0 {
            // Every sample coincides with a centre — a perfectly flat side.
            // More modes would all be the same mode.
            break;
        }

        let mut target = rng.unit() as f64 * total;
        let mut pick = samples.len() - 1;
        for (s, &d) in nearest.iter().enumerate() {
            target -= d as f64;
            if target <= 0.0 {
                pick = s;
                break;
            }
        }
        centres.push(samples[pick]);
    }
    centres
}

/// The closest mode, as `(squared Mahalanobis, negative log density)`, both
/// charged for that mode's rarity.
///
/// The rarity enters the Mahalanobis term doubled because that term is a
/// squared distance and the other is a log density: a chi-square of `2r` is
/// the same amount of evidence as `r` nats, so the two tests are then denominated
/// in the same currency and the strictness means one thing across both.
///
/// Nearest rather than a proper mixture sum: the largest term dominates a
/// well-separated mixture, and taking the max is what makes the two numbers
/// this returns describe *the same* mode.
fn nearest(modes: &[Mode], x: &[f32; FEATURES]) -> (f32, f32) {
    let mut best = (f32::INFINITY, f32::INFINITY);
    for mode in modes {
        let mut maha = 0.0f32;
        for ((v, mean), variance) in x.iter().zip(&mode.mean).zip(&mode.variance) {
            let d = v - mean;
            maha += d * d / variance;
        }
        let nll = 0.5 * maha + mode.half_log_det + mode.rarity;
        if nll < best.1 {
            best = (maha + 2.0 * mode.rarity, nll);
        }
    }
    best
}

fn squared(a: &[f32; FEATURES], b: &[f32; FEATURES]) -> f32 {
    (0..FEATURES).map(|f| (a[f] - b[f]).powi(2)).sum()
}

fn smoothstep(edge0: f32, edge1: f32, x: f32) -> f32 {
    if edge1 <= edge0 {
        return if x >= edge1 { 1.0 } else { 0.0 };
    }
    let t = ((x - edge0) / (edge1 - edge0)).clamp(0.0, 1.0);
    t * t * (3.0 - 2.0 * t)
}

/// Separable box blur, clamped at the edges.
///
/// A box rather than a gaussian because what is wanted is the removal of
/// single-pixel speckle, and the shape of the kernel that does it does not
/// matter.
fn blur(src: &[f32], width: usize, height: usize, radius: usize) -> Vec<f32> {
    if radius == 0 || width == 0 || height == 0 {
        return src.to_vec();
    }
    let span = (2 * radius + 1) as f32;
    let mut mid = vec![0.0f32; src.len()];
    for y in 0..height {
        for x in 0..width {
            let mut sum = 0.0f32;
            for k in 0..=2 * radius {
                let sx = (x + k).saturating_sub(radius).min(width - 1);
                sum += src[y * width + sx];
            }
            mid[y * width + x] = sum / span;
        }
    }
    let mut out = vec![0.0f32; src.len()];
    for y in 0..height {
        for x in 0..width {
            let mut sum = 0.0f32;
            for k in 0..=2 * radius {
                let sy = (y + k).saturating_sub(radius).min(height - 1);
                sum += mid[sy * width + x];
            }
            out[y * width + x] = sum / span;
        }
    }
    out
}

/// A fixed-sequence generator, for the one place a draw is needed.
struct Lcg(u64);

impl Lcg {
    fn unit(&mut self) -> f32 {
        self.0 = self
            .0
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        // The high bits are the well-mixed ones in an LCG; the low bits cycle
        // with a short period and would make the draw far from uniform.
        ((self.0 >> 40) as f32) / ((1u64 << 24) as f32)
    }
}

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

    /// One logit cell, in the pixels of the synthetic frames below.
    const CELL: f32 = 8.0;

    /// Edge of every synthetic frame.
    ///
    /// Large enough that an intruder can be a realistic *fraction* of the
    /// category rather than a realistic number of pixels — which is what
    /// rarity is measured in, and getting that proportion wrong is the
    /// difference between this method working and not.
    const EDGE: usize = 192;

    /// A blue upper half over a green lower half, and a category mask that
    /// claims the whole upper half — which is what the coarse grid produces.
    fn landscape(w: usize, h: usize) -> (Vec<f32>, Vec<f32>) {
        let mut rgb = vec![0.0f32; w * h * 3];
        let mut weights = vec![0.0f32; w * h];
        for y in 0..h {
            for x in 0..w {
                let p = y * w + x;
                // A vertical gradient in the sky, because a real one has one
                // and a single Gaussian over luma is exactly what it breaks.
                let t = y as f32 / h as f32;
                if y < h / 2 {
                    rgb[p * 3] = 0.25 + 0.4 * t;
                    rgb[p * 3 + 1] = 0.45 + 0.35 * t;
                    rgb[p * 3 + 2] = 0.85;
                } else {
                    rgb[p * 3] = 0.20;
                    rgb[p * 3 + 1] = 0.45;
                    rgb[p * 3 + 2] = 0.15;
                }
                weights[p] = if y < h / 2 { 1.0 } else { 0.0 };
            }
        }
        (rgb, weights)
    }

    /// Paint a rectangle into the frame, and let the coarse mask claim it —
    /// the flag in the sky.
    fn intrude(
        rgb: &mut [f32],
        weights: &mut [f32],
        w: usize,
        rect: (usize, usize, usize, usize),
        colour: [f32; 3],
    ) {
        let (x0, y0, x1, y1) = rect;
        for y in y0..y1 {
            for x in x0..x1 {
                let p = y * w + x;
                rgb[p * 3..p * 3 + 3].copy_from_slice(&colour);
                weights[p] = 1.0;
            }
        }
    }

    fn mean(v: &[f32], w: usize, rect: (usize, usize, usize, usize)) -> f32 {
        let (x0, y0, x1, y1) = rect;
        let mut sum = 0.0;
        for y in y0..y1 {
            for x in x0..x1 {
                sum += v[y * w + x];
            }
        }
        sum / ((x1 - x0) * (y1 - y0)) as f32
    }

    const FLAG: (usize, usize, usize, usize) = (88, 40, 104, 56);
    const FLAG_CORE: (usize, usize, usize, usize) = (92, 44, 100, 52);
    const CLOUD: (usize, usize, usize, usize) = (30, 12, 150, 60);
    const CLOUD_CORE: (usize, usize, usize, usize) = (50, 24, 130, 48);
    const RED: [f32; 3] = [0.75, 0.10, 0.12];
    const WHITE: [f32; 3] = [0.92, 0.94, 0.96];

    fn with(rect: (usize, usize, usize, usize), colour: [f32; 3]) -> (Vec<f32>, Vec<f32>) {
        let (mut rgb, mut weights) = landscape(EDGE, EDGE);
        intrude(&mut rgb, &mut weights, EDGE, rect, colour);
        (rgb, weights)
    }

    /// The case the module exists for: a red flag inside the sky, claimed by
    /// the coarse mask, must come back out at the default strictness — while
    /// the sky around it stays.
    #[test]
    fn a_flag_in_the_sky_is_removed() {
        let (rgb, weights) = with(FLAG, RED);
        let (out, what) = refine_category(
            &weights,
            &rgb,
            EDGE,
            EDGE,
            CELL,
            STRICTNESS_DEFAULT,
            &RefineOptions::default(),
        );
        assert!(
            matches!(what, Refined::Applied { .. }),
            "should have had the evidence to judge: {what:?}"
        );

        let inside = mean(&out, EDGE, FLAG_CORE);
        assert!(inside < 0.2, "the flag should be cut out, got {inside}");

        let sky = mean(&out, EDGE, (8, 8, 60, 60));
        assert!(sky > 0.8, "the sky around it should survive, got {sky}");
    }

    /// A cloud is a legitimate part of the sky and is separated from the blue
    /// by luminance alone, which is what `luma_weight` is for — and it holds a
    /// large share, which is what rarity is for.
    #[test]
    fn a_cloud_is_not_mistaken_for_an_intruder() {
        let (rgb, weights) = with(CLOUD, WHITE);
        let (out, _) = refine_category(
            &weights,
            &rgb,
            EDGE,
            EDGE,
            CELL,
            STRICTNESS_DEFAULT,
            &RefineOptions::default(),
        );
        let cloud = mean(&out, EDGE, CLOUD_CORE);
        assert!(cloud > 0.7, "the cloud should stay sky, got {cloud}");
    }

    /// The property that makes a single slider worth offering.
    ///
    /// Rarity puts the flag and the cloud at *different places on one scale*
    /// rather than on two sides of a threshold, so there is a strictness that
    /// has taken the flag and not the cloud — and the flag always goes first.
    /// Without the rarity term both sit at the same verdict and the control
    /// would appear dead until it suddenly ate the sky.
    #[test]
    fn the_flag_goes_before_the_cloud_does() {
        let opts = RefineOptions::default();

        let (flag_rgb, flag_w) = with(FLAG, RED);
        let flag = Refinement::compute(&flag_w, &flag_rgb, EDGE, EDGE, CELL, &opts)
            .expect("the flag frame has both sides");

        let (cloud_rgb, cloud_w) = with(CLOUD, WHITE);
        let cloud = Refinement::compute(&cloud_w, &cloud_rgb, EDGE, EDGE, CELL, &opts)
            .expect("the cloud frame has both sides");

        // The first strictness on a fine sweep at which each is more than half
        // gone. `None` would mean it never goes at all.
        let crossing = |r: &Refinement, weights: &[f32], core| {
            (0..=120)
                .map(|i| i as f32 * STRICTNESS_MAX / 120.0)
                .find(|&s| mean(&r.apply(weights, s), EDGE, core) < 0.5)
        };

        let flag_at = crossing(&flag, &flag_w, FLAG_CORE).expect("the flag must go somewhere");
        let cloud_at = crossing(&cloud, &cloud_w, CLOUD_CORE);

        // `None` is the better outcome, not a missing case: it means the cloud
        // survived the whole range, so the flag went first by a margin wider
        // than the control can travel.
        if let Some(cloud_at) = cloud_at {
            assert!(
                flag_at < cloud_at,
                "the flag must go first: flag at {flag_at}, cloud at {cloud_at}"
            );
        }
        assert!(
            flag_at <= STRICTNESS_DEFAULT,
            "the default must already clear a flag, but it only goes at {flag_at}"
        );
    }

    /// Zero strictness is exactly the model's own weighting.
    ///
    /// The control's off position has to be the old behaviour bit for bit, or
    /// "turn it off and compare" does not answer the question it is asked.
    #[test]
    fn strictness_zero_changes_nothing() {
        let (rgb, weights) = with(FLAG, RED);
        let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
            .expect("both sides present");
        assert_eq!(r.apply(&weights, STRICTNESS_OFF), weights);
    }

    /// Raising the strictness may only ever remove more.
    ///
    /// A slider that gave weight back on the way up would be one whose
    /// direction the user cannot predict, and it would break the ordering the
    /// whole control rests on.
    #[test]
    fn strictness_is_monotonic() {
        let (rgb, weights) = with(FLAG, RED);
        let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
            .expect("both sides present");

        let mut previous = r.apply(&weights, 0.0);
        for step in 1..=12 {
            let next = r.apply(&weights, step as f32 * STRICTNESS_MAX / 12.0);
            for (p, (&before, &after)) in previous.iter().zip(&next).enumerate() {
                assert!(
                    after <= before + 1e-6,
                    "pixel {p} gained weight at step {step}: {before} -> {after}"
                );
            }
            previous = next;
        }
    }

    /// The refinement may only ever take weight away.
    ///
    /// This is what bounds its damage and what keeps `scene.rs`'s partition
    /// true when several categories are refined independently, so it is
    /// asserted rather than left as a property of the arithmetic.
    #[test]
    fn refinement_is_subtractive() {
        let (rgb, weights) = with(FLAG, RED);
        let (out, _) = refine_category(
            &weights,
            &rgb,
            EDGE,
            EDGE,
            CELL,
            STRICTNESS_DEFAULT,
            &RefineOptions::default(),
        );
        for (p, (&before, &after)) in weights.iter().zip(&out).enumerate() {
            assert!(
                after <= before + 1e-6,
                "pixel {p} gained weight: {before} -> {after}"
            );
            assert!(
                (0.0..=1.0).contains(&after),
                "pixel {p} out of range: {after}"
            );
        }
    }

    /// The coverage path must agree with the float one, since the application
    /// uses the first and every test here uses the second.
    #[test]
    fn the_coverage_path_agrees_with_the_float_one() {
        let (rgb, weights) = with(FLAG, RED);
        let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
            .expect("both sides present");

        let coverage: Vec<u8> = weights.iter().map(|&w| (w * 255.0).round() as u8).collect();
        let bytes = r.apply_coverage(&coverage, STRICTNESS_DEFAULT);
        let floats = r.apply(&weights, STRICTNESS_DEFAULT);

        for (p, (&b, &f)) in bytes.iter().zip(&floats).enumerate() {
            let expected = (f * 255.0).round() as u8;
            assert!(
                b.abs_diff(expected) <= 1,
                "pixel {p}: coverage {b}, float {expected}"
            );
        }
    }

    /// No confident interior means no distribution, and no distribution must
    /// mean "leave it alone" rather than "reject everything".
    #[test]
    fn a_mask_thinner_than_the_grid_is_left_alone() {
        let (rgb, _) = landscape(EDGE, EDGE);
        // A three-pixel stripe: narrower than one cell, so the erosion at
        // 1.5 cells empties it.
        let mut weights = vec![0.0f32; EDGE * EDGE];
        for y in 0..EDGE {
            for x in 60..63 {
                weights[y * EDGE + x] = 1.0;
            }
        }

        assert_eq!(
            Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default()),
            Err(SkipReason::NoInterior)
        );
        let (out, what) = refine_category(
            &weights,
            &rgb,
            EDGE,
            EDGE,
            CELL,
            STRICTNESS_DEFAULT,
            &RefineOptions::default(),
        );
        assert_eq!(what, Refined::Skipped(SkipReason::NoInterior));
        assert_eq!(out, weights, "a skip must return the input untouched");
    }

    /// A frame that is entirely one category has nothing to contrast against.
    #[test]
    fn a_frame_of_nothing_but_sky_is_left_alone() {
        let rgb = vec![0.6f32; EDGE * EDGE * 3];
        let weights = vec![1.0f32; EDGE * EDGE];

        let (out, what) = refine_category(
            &weights,
            &rgb,
            EDGE,
            EDGE,
            CELL,
            STRICTNESS_DEFAULT,
            &RefineOptions::default(),
        );
        assert_eq!(what, Refined::Skipped(SkipReason::NoExterior));
        assert_eq!(out, weights);
    }

    /// Masks reach the sidecar as indices, so the same input must give the
    /// same mask on every run and every machine — which is why the k-means
    /// seeding uses a fixed sequence and the iteration count is fixed.
    #[test]
    fn the_same_input_gives_the_same_answer() {
        let (rgb, weights) = with(FLAG, RED);
        let opts = RefineOptions::default();
        let a = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &opts).unwrap();
        let b = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &opts).unwrap();
        assert_eq!(
            a.apply(&weights, STRICTNESS_DEFAULT),
            b.apply(&weights, STRICTNESS_DEFAULT)
        );
    }

    #[test]
    fn buffers_that_disagree_are_refused() {
        let (out, what) = refine_category(
            &[1.0; 4],
            &[0.5; 6],
            2,
            2,
            CELL,
            STRICTNESS_DEFAULT,
            &RefineOptions::default(),
        );
        assert_eq!(what, Refined::Skipped(SkipReason::Mismatched));
        assert_eq!(out, vec![1.0; 4]);
    }
}