visual_cortex_vision/

stability_mask.rs

//! Block-state core for the temporal stability mask. Pure grid math —
//! no I/O, no frames — so every rule is unit-testable in isolation.

use std::sync::Arc;
use std::time::Duration;

/// Per-block temporal state. Signatures are mean-luma u8.
#[derive(Debug, Clone)]
pub(crate) struct BlockState {
    pub last_sig: u8,
    /// Consecutive ticks the signature stayed within tolerance.
    pub stable_ticks: u32,
    /// Slow-moving estimate of the block's "permanent" appearance.
    pub baseline: f32,
    /// Set while the block is kept: the baseline stops absorbing (a held
    /// tooltip must not become baseline and vanish mid-hover).
    pub baseline_frozen: bool,
    pub initialized: bool,
}

impl Default for BlockState {
    fn default() -> Self {
        Self {
            last_sig: 0,
            stable_ticks: 0,
            baseline: 0.0,
            baseline_frozen: false,
            initialized: false,
        }
    }
}

pub(crate) struct MaskParams {
    pub stable_ticks: u32,
    /// EMA horizon in ticks (alpha = 1/horizon).
    pub baseline_ticks: u32,
    pub signature_tolerance: u8,
    pub baseline_threshold: u8,
}

/// Advance one block one tick. Returns whether the block is KEPT this tick
/// (stable AND novel vs baseline), updating all state per the spec:
/// - stability counter resets on movement;
/// - baseline is EMA-updated only while NOT frozen;
/// - freeze begins when the block is kept, and thaws only when the block is
///   no longer kept AND its signature has returned near the frozen baseline.
pub(crate) fn step_block(state: &mut BlockState, sig: u8, p: &MaskParams) -> bool {
    if !state.initialized {
        *state = BlockState {
            last_sig: sig,
            stable_ticks: 0,
            baseline: sig as f32,
            baseline_frozen: false,
            initialized: true,
        };
        return false; // cold start: nothing is stable yet
    }

    // Stability.
    if sig.abs_diff(state.last_sig) <= p.signature_tolerance {
        state.stable_ticks = state.stable_ticks.saturating_add(1);
    } else {
        state.stable_ticks = 0;
    }
    state.last_sig = sig;

    let stable = state.stable_ticks >= p.stable_ticks;
    let novel = (sig as f32 - state.baseline).abs() > p.baseline_threshold as f32;
    let kept = stable && novel;

    // Baseline bookkeeping.
    if kept {
        state.baseline_frozen = true;
    } else if state.baseline_frozen {
        // Thaw only once the overlay has actually gone: signature back near
        // the frozen baseline.
        if (sig as f32 - state.baseline).abs() <= p.baseline_threshold as f32 {
            state.baseline_frozen = false;
        }
    }
    if !state.baseline_frozen {
        let alpha = 1.0 / p.baseline_ticks.max(1) as f32;
        state.baseline += alpha * (sig as f32 - state.baseline);
    }

    kept
}

/// Morphological close (fill interior holes) then dilate by `dilate` rings,
/// on a row-major block grid. Close = dilate-then-erode with a 1-ring
/// structuring element.
pub(crate) fn close_and_dilate(keep: &[bool], cols: usize, rows: usize, dilate: u32) -> Vec<bool> {
    let d1 = dilate_grid(keep, cols, rows, 1);
    let closed = erode_grid(&d1, cols, rows, 1);
    if dilate == 0 {
        closed
    } else {
        dilate_grid(&closed, cols, rows, dilate as usize)
    }
}

fn dilate_grid(grid: &[bool], cols: usize, rows: usize, rings: usize) -> Vec<bool> {
    let mut out = grid.to_vec();
    for _ in 0..rings {
        let src = out.clone();
        for r in 0..rows {
            for c in 0..cols {
                if src[r * cols + c] {
                    continue;
                }
                let neighbors_on = (r > 0 && src[(r - 1) * cols + c])
                    || (r + 1 < rows && src[(r + 1) * cols + c])
                    || (c > 0 && src[r * cols + c - 1])
                    || (c + 1 < cols && src[r * cols + c + 1]);
                if neighbors_on {
                    out[r * cols + c] = true;
                }
            }
        }
    }
    out
}

fn erode_grid(grid: &[bool], cols: usize, rows: usize, rings: usize) -> Vec<bool> {
    let mut out = grid.to_vec();
    for _ in 0..rings {
        let src = out.clone();
        for r in 0..rows {
            for c in 0..cols {
                if !src[r * cols + c] {
                    continue;
                }
                let all_on = (r == 0 || src[(r - 1) * cols + c])
                    && (r + 1 >= rows || src[(r + 1) * cols + c])
                    && (c == 0 || src[r * cols + c - 1])
                    && (c + 1 >= cols || src[r * cols + c + 1]);
                if !all_on {
                    out[r * cols + c] = false;
                }
            }
        }
    }
    out
}

use visual_cortex_capture::{Frame, FrameView, Rate};

use crate::debug::{DebugSink, DebugStage};
use crate::detector::DetectorError;
use crate::preprocessor::Preprocessor;

/// Temporal stability mask: keeps only blocks that are *stable* over recent
/// ticks yet *novel* versus a longer-term baseline — isolating
/// newly-appeared static overlays (tooltips) from moving gameplay and
/// permanent HUD. Suppressed blocks render black; kept blocks pass through
/// byte-identical. See the design spec for the algebra.
pub struct StabilityMask {
    block: u32,
    params: MaskParams,
    dilate: u32,
    /// Grid state; reset when the view dimensions change.
    state: Vec<BlockState>,
    dims: (u32, u32),
    tick: u64,
    sinks: Vec<Box<dyn DebugSink>>,
}

impl Default for StabilityMask {
    fn default() -> Self {
        Self::new()
    }
}

impl StabilityMask {
    pub fn new() -> Self {
        Self {
            block: 16,
            params: MaskParams {
                stable_ticks: 6,
                baseline_ticks: 100,
                signature_tolerance: 4,
                baseline_threshold: 25,
            },
            dilate: 1,
            state: Vec::new(),
            dims: (0, 0),
            tick: 0,
            sinks: Vec::new(),
        }
    }

    /// Analysis block size in pixels (default 16).
    pub fn block_size(mut self, px: u32) -> Self {
        assert!(px > 0, "block size must be positive");
        self.block = px;
        self
    }

    /// Consecutive stable ticks before a block counts as stable (default 6).
    pub fn stable_ticks(mut self, k: u32) -> Self {
        self.params.stable_ticks = k;
        self
    }

    /// Duration form of [`Self::stable_ticks`]; converted via the watcher's
    /// rate (tick-based internals keep paused-time tests deterministic).
    pub fn stable_for(self, rate: Rate, duration: Duration) -> Self {
        let ticks = (duration.as_secs_f64() / rate.period().as_secs_f64()).ceil() as u32;
        self.stable_ticks(ticks.max(1))
    }

    /// Baseline EMA horizon in ticks (default 100).
    pub fn baseline_ticks(mut self, t: u32) -> Self {
        self.params.baseline_ticks = t.max(1);
        self
    }

    /// Duration form of [`Self::baseline_ticks`].
    pub fn baseline(self, rate: Rate, duration: Duration) -> Self {
        let ticks = (duration.as_secs_f64() / rate.period().as_secs_f64()).ceil() as u32;
        self.baseline_ticks(ticks.max(1))
    }

    /// Keep-mask dilation rings (default 1) — keeps glyphs off the fill
    /// boundary and closes holes inside kept clusters.
    pub fn dilate(mut self, rings: u32) -> Self {
        self.dilate = rings;
        self
    }

    /// Stability comparison tolerance on the 0-255 luma signature (default 4).
    pub fn signature_tolerance(mut self, tol: u8) -> Self {
        self.params.signature_tolerance = tol;
        self
    }

    /// Novelty threshold versus the baseline (default 25).
    pub fn baseline_threshold(mut self, thr: u8) -> Self {
        self.params.baseline_threshold = thr;
        self
    }

    /// Attach a debug tap (e.g. [`PngDump`](crate::PngDump)) receiving
    /// per-tick Input/Baseline/Overlay/Output frames. Composable —
    /// multiple sinks allowed. Stage frames are only materialized when at
    /// least one sink is attached.
    pub fn debug_sink(mut self, sink: impl DebugSink) -> Self {
        self.sinks.push(Box::new(sink));
        self
    }

    fn grid(&self, w: u32, h: u32) -> (usize, usize) {
        (
            (w as usize).div_ceil(self.block as usize),
            (h as usize).div_ceil(self.block as usize),
        )
    }
}

/// Mean luma per block, row-major grid. Blocks at the right/bottom edges may
/// be partial.
fn block_signatures(view: &FrameView<'_>, block: u32, cols: usize, rows: usize) -> Vec<u8> {
    let mut sums = vec![0u64; cols * rows];
    let mut counts = vec![0u64; cols * rows];
    for (y, row) in view.rows().enumerate() {
        let br = y / block as usize;
        for (x, px) in row.chunks_exact(4).enumerate() {
            let bc = x / block as usize;
            // Integer BT.601 luma, same weights as the detectors.
            let luma = (px[2] as u32 * 299 + px[1] as u32 * 587 + px[0] as u32 * 114) / 1000;
            sums[br * cols + bc] += luma as u64;
            counts[br * cols + bc] += 1;
        }
    }
    sums.iter()
        .zip(&counts)
        .map(|(s, c)| if *c == 0 { 0 } else { (s / c) as u8 })
        .collect()
}

impl Preprocessor for StabilityMask {
    fn process(&mut self, view: &FrameView<'_>) -> Result<Arc<Frame>, DetectorError> {
        let (w, h) = (view.width(), view.height());
        let (cols, rows) = self.grid(w, h);
        if self.dims != (w, h) {
            self.state = vec![BlockState::default(); cols * rows];
            self.dims = (w, h);
        }

        let sigs = block_signatures(view, self.block, cols, rows);
        let keep_raw: Vec<bool> = sigs
            .iter()
            .zip(self.state.iter_mut())
            .map(|(sig, st)| step_block(st, *sig, &self.params))
            .collect();
        let keep = close_and_dilate(&keep_raw, cols, rows, self.dilate);

        // Render: kept blocks byte-identical, suppressed blocks black.
        let mut data = vec![0u8; w as usize * h as usize * 4];
        for px in data.chunks_exact_mut(4) {
            px[3] = 255;
        }
        for (y, row) in view.rows().enumerate() {
            let br = y / self.block as usize;
            for bc in 0..cols {
                if !keep[br * cols + bc] {
                    continue;
                }
                let x0 = bc * self.block as usize * 4;
                let x1 = (((bc + 1) * self.block as usize) * 4).min(row.len());
                let dst = y * w as usize * 4;
                data[dst + x0..dst + x1].copy_from_slice(&row[x0..x1]);
            }
        }
        let output = Frame::new(w, h, data)
            .map(Arc::new)
            .map_err(|e| DetectorError::Other(format!("mask render: {e}")))?;

        if !self.sinks.is_empty() {
            self.emit_debug(view, &keep, cols, &output);
        }
        self.tick += 1;
        Ok(output)
    }

    fn set_label(&mut self, label: &str) {
        for sink in &mut self.sinks {
            sink.set_label(label);
        }
    }
}

impl StabilityMask {
    /// Materialize the debug stage frames and fan them out. Only called
    /// when sinks exist; sinks themselves must never block (drop-on-lag).
    fn emit_debug(&mut self, view: &FrameView<'_>, keep: &[bool], cols: usize, output: &Frame) {
        let (w, h) = (view.width(), view.height());
        let input = Frame::new(w, h, view.to_vec()).expect("view-sized buffer");

        // Baseline: each block filled with its EMA estimate, as gray.
        // Overlay: input with suppressed blocks dimmed to 25% — the human view.
        let mut baseline = vec![0u8; w as usize * h as usize * 4];
        let mut overlay = input.data().to_vec();
        for y in 0..h as usize {
            let br = y / self.block as usize;
            for x in 0..w as usize {
                let bc = x / self.block as usize;
                let i = (y * w as usize + x) * 4;
                let b = self.state[br * cols + bc].baseline as u8;
                baseline[i] = b;
                baseline[i + 1] = b;
                baseline[i + 2] = b;
                baseline[i + 3] = 255;
                if !keep[br * cols + bc] {
                    overlay[i] /= 4;
                    overlay[i + 1] /= 4;
                    overlay[i + 2] /= 4;
                }
            }
        }
        let baseline = Frame::new(w, h, baseline).expect("view-sized buffer");
        let overlay = Frame::new(w, h, overlay).expect("view-sized buffer");

        for sink in &mut self.sinks {
            sink.write(self.tick, DebugStage::Input, &input);
            sink.write(self.tick, DebugStage::Baseline, &baseline);
            sink.write(self.tick, DebugStage::Overlay, &overlay);
            sink.write(self.tick, DebugStage::Output, output);
        }
    }
}

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

    fn params() -> MaskParams {
        MaskParams {
            stable_ticks: 3,
            baseline_ticks: 20,
            signature_tolerance: 4,
            baseline_threshold: 25,
        }
    }

    fn run(state: &mut BlockState, sigs: &[u8], p: &MaskParams) -> Vec<bool> {
        sigs.iter().map(|s| step_block(state, *s, p)).collect()
    }

    #[test]
    fn constant_hud_block_is_never_kept() {
        // Same signature forever: stable but never novel (baseline == sig).
        let mut st = BlockState::default();
        let kept = run(&mut st, &[100; 30], &params());
        assert!(kept.iter().all(|k| !k), "HUD must stay suppressed");
    }

    #[test]
    fn noisy_gameplay_block_is_never_kept() {
        // Signature jumps beyond tolerance every tick: never stable.
        let mut st = BlockState::default();
        let sigs: Vec<u8> = (0..30).map(|i| if i % 2 == 0 { 40 } else { 200 }).collect();
        let kept = run(&mut st, &sigs, &params());
        assert!(
            kept.iter().all(|k| !k),
            "moving content must stay suppressed"
        );
    }

    #[test]
    fn tooltip_becomes_kept_after_k_stable_ticks() {
        let mut st = BlockState::default();
        // Settle on gameplay-ish baseline, then a tooltip appears (sig 200).
        run(&mut st, &[60; 10], &params());
        let kept = run(&mut st, &[200; 6], &params());
        // Tick 1 resets stability (jump); the counter then reaches K=3 on the
        // fourth tick of the new signature (index 3).
        assert_eq!(kept, vec![false, false, false, true, true, true]);
    }

    #[test]
    fn held_tooltip_survives_past_the_baseline_horizon() {
        // The ghosting trap: without baseline freeze, ~baseline_ticks of a
        // held tooltip would absorb it into the baseline and un-keep it.
        let mut st = BlockState::default();
        run(&mut st, &[60; 10], &params());
        let long_hold = run(&mut st, &[200; 200], &params()); // 10x horizon
        assert!(
            long_hold[10..].iter().all(|k| *k),
            "held tooltip must remain kept indefinitely"
        );
    }

    #[test]
    fn baseline_thaws_after_overlay_departs() {
        let mut st = BlockState::default();
        run(&mut st, &[60; 10], &params());
        run(&mut st, &[200; 20], &params()); // tooltip held (frozen baseline)
        assert!(st.baseline_frozen);
        // Tooltip dismissed: back to the old scene.
        let after = run(&mut st, &[60; 10], &params());
        assert!(!st.baseline_frozen, "baseline resumes after departure");
        assert!(after.iter().all(|k| !k), "restored scene is not novel");
    }

    #[test]
    fn close_fills_interior_holes_and_dilate_expands() {
        // 5x5 grid: ring of kept blocks with a hole in the middle.
        let cols = 5;
        let mut keep = vec![false; 25];
        for (r, c) in [
            (1usize, 1usize),
            (1, 2),
            (1, 3),
            (2, 1),
            (2, 3),
            (3, 1),
            (3, 2),
            (3, 3),
        ] {
            keep[r * cols + c] = true;
        }
        let out = close_and_dilate(&keep, cols, 5, 0);
        assert!(out[2 * cols + 2], "interior hole must be closed");
        let out = close_and_dilate(&keep, cols, 5, 1);
        assert!(out[cols], "edge cell (1,0) reached by one dilation ring");
        assert!(
            !out[0],
            "corner is two steps away; one ring must not reach it"
        );
    }

    #[test]
    fn debug_sinks_receive_all_four_stages_per_tick() {
        use std::sync::Mutex;

        use crate::debug::{DebugSink, DebugStage};
        use visual_cortex_capture::PxRect;

        struct Recording(Arc<Mutex<Vec<(u64, DebugStage)>>>);
        impl DebugSink for Recording {
            fn write(&mut self, tick: u64, stage: DebugStage, _frame: &Frame) {
                self.0.lock().unwrap().push((tick, stage));
            }
        }

        let log = Arc::new(Mutex::new(Vec::new()));
        let mut mask = StabilityMask::new()
            .block_size(8)
            .debug_sink(Recording(log.clone()));
        let frame = Frame::solid(8, 8, [60, 60, 60, 255]);
        let view = frame
            .view(PxRect {
                x: 0,
                y: 0,
                w: 8,
                h: 8,
            })
            .unwrap();
        mask.process(&view).unwrap();
        mask.process(&view).unwrap();

        use DebugStage::*;
        assert_eq!(
            *log.lock().unwrap(),
            vec![
                (0, Input),
                (0, Baseline),
                (0, Overlay),
                (0, Output),
                (1, Input),
                (1, Baseline),
                (1, Overlay),
                (1, Output),
            ]
        );
    }

    #[test]
    fn cold_start_keeps_nothing() {
        let mut st = BlockState::default();
        assert!(!step_block(&mut st, 200, &params()));
    }
}