darkly 0.5.0

//! Laplacian relaxation stabilizer — iterative smoothing with zero lag.
//!
//! Maintains a polyline of all raw input positions.  On each new input:
//! 1. Append to polyline
//! 2. Run N iterations of Laplacian smoothing on interior points (first + last pinned)
//! 3. Each iteration: `point[i] = lerp(point[i], avg(point[i-1], point[i+1]), strength)`
//! 4. Sensor values (pressure, tilt, etc.) smoothed the same way
//! 5. Diff against previous frame's polyline → find divergence point
//!
//! The tip is always pinned at the cursor (zero lag).  The stroke behind
//! the pen continuously reshapes as direction changes — the "taffy" feel.

use crate::brush::paint_info::PaintInformation;
use crate::brush::stabilizer::{StabilizeResult, StabilizerAlgorithm, StabilizerRegistration};
use crate::gpu::params::{ParamDef, ParamValue};

/// Threshold in pixels below which a point is considered unchanged.
const DIVERGENCE_EPSILON: f32 = 0.5;

const PARAMS: &[ParamDef] = &[ParamDef::Float {
    name: "strength",
    min: 0.0,
    max: 1.0,
    default: 0.5,
}];

pub fn register() -> StabilizerRegistration {
    StabilizerRegistration {
        type_id: "laplacian",
        display_name: "Laplacian Relaxation",
        params: PARAMS,
        from_params: |params| {
            let strength = match params.first() {
                Some(ParamValue::Float(v)) => *v,
                _ => 0.5,
            };
            Box::new(LaplacianStabilizer::new(strength))
        },
    }
}

pub struct LaplacianStabilizer {
    raw_points: Vec<PaintInformation>,
    stabilized: Vec<PaintInformation>,
    prev_positions: Vec<[f32; 2]>,
    strength: f32,
}

impl LaplacianStabilizer {
    pub fn new(strength: f32) -> Self {
        Self {
            raw_points: Vec::with_capacity(256),
            stabilized: Vec::with_capacity(256),
            prev_positions: Vec::with_capacity(256),
            strength: strength.clamp(0.0, 1.0),
        }
    }

    /// Run Laplacian relaxation on the stabilized polyline.
    /// First and last points are pinned (never move).
    fn relax(&mut self) {
        let len = self.stabilized.len();
        if len < 3 {
            return;
        }

        let iterations = (self.strength * 10.0).ceil() as u32;
        let s = self.strength;

        for _ in 0..iterations {
            for i in 1..len - 1 {
                // Position smoothing.
                let prev_pos = self.stabilized[i - 1].pos;
                let next_pos = self.stabilized[i + 1].pos;
                let avg = [
                    (prev_pos[0] + next_pos[0]) * 0.5,
                    (prev_pos[1] + next_pos[1]) * 0.5,
                ];
                let cur = &mut self.stabilized[i];
                cur.pos[0] += (avg[0] - cur.pos[0]) * s;
                cur.pos[1] += (avg[1] - cur.pos[1]) * s;

                // Sensor smoothing — same treatment for all continuous values.
                let prev = self.stabilized[i - 1];
                let next = self.stabilized[i + 1];
                let cur = &mut self.stabilized[i];

                macro_rules! smooth_field {
                    ($field:ident) => {
                        let avg = (prev.$field + next.$field) * 0.5;
                        cur.$field += (avg - cur.$field) * s;
                    };
                }

                smooth_field!(pressure);
                smooth_field!(x_tilt);
                smooth_field!(y_tilt);
                smooth_field!(rotation);
                smooth_field!(tangential_pressure);
                smooth_field!(speed);
                smooth_field!(tilt_magnitude);
                smooth_field!(tilt_direction);
            }
        }
    }

    /// Find the divergence point: walk backward from tip until either the
    /// position delta between current and previous frame falls below
    /// `DIVERGENCE_EPSILON`, or we hit the earliest index the influence model
    /// admits a perturbation could reach.
    ///
    /// The walk is bounded by [`Self::max_divergence_window`] — both methods
    /// share the same Laplacian-relaxation influence model, so the bound is
    /// enforced by construction rather than by clamping a wider scan. See
    /// `max_divergence_window` for the derivation.
    fn find_divergence(&self) -> Option<usize> {
        let len = self.stabilized.len();
        if len == 0 {
            return None;
        }

        // `earliest` is the lowest index whose stabilized position could
        // possibly differ from the previous frame given the influence radius.
        // Walking past it is wasted work and risks reporting spurious
        // divergence at indices the model says cannot have moved.
        let max_back = self.max_divergence_window();
        let earliest = len.saturating_sub(max_back + 1);
        let eps2 = DIVERGENCE_EPSILON * DIVERGENCE_EPSILON;

        // Helper: squared position delta at index `i` (always within bounds
        // of both arrays here — callers ensure that).
        let delta2 = |i: usize| -> f32 {
            let cur = self.stabilized[i].pos;
            let prev = self.prev_positions[i];
            let dx = cur[0] - prev[0];
            let dy = cur[1] - prev[1];
            dx * dx + dy * dy
        };

        if self.prev_positions.len() == len {
            // Same length — walk backward from tip to `earliest`.
            for i in (earliest..len).rev() {
                if delta2(i) < eps2 {
                    return if i + 1 < len { Some(i + 1) } else { None };
                }
            }
            Some(earliest)
        } else {
            // Polyline grew (typically by exactly one). New indices
            // `[prev_positions.len(), len-1]` are by definition new and
            // cannot be compared. Existing indices that overlap with
            // `prev_positions` are in `[0, overlap_end)` — walk those,
            // descending, bounded below by `earliest`.
            let overlap_end = self.prev_positions.len().min(len);
            if earliest >= overlap_end {
                // No overlap to check (e.g., first push of a stroke). The
                // divergence index is `earliest` itself, which equals 0
                // when there is nothing prior to compare against.
                return Some(earliest);
            }
            for i in (earliest..overlap_end).rev() {
                if delta2(i) < eps2 {
                    return Some(i + 1);
                }
            }
            Some(earliest)
        }
    }
}

impl StabilizerAlgorithm for LaplacianStabilizer {
    fn push(&mut self, point: PaintInformation) -> StabilizeResult {
        // Save previous positions for divergence detection.
        self.prev_positions.clear();
        self.prev_positions
            .extend(self.stabilized.iter().map(|p| p.pos));

        // Append raw point.
        self.raw_points.push(point);

        // Copy raw → stabilized (fresh copy each frame for correct relaxation).
        self.stabilized.clear();
        self.stabilized.extend_from_slice(&self.raw_points);

        // Run relaxation.
        self.relax();

        // Find divergence.
        let divergence_index = if self.strength == 0.0 {
            None
        } else {
            self.find_divergence()
        };

        StabilizeResult { divergence_index }
    }

    fn stabilized(&self) -> &[PaintInformation] {
        &self.stabilized
    }

    /// Conservative upper bound on `tip_vi - find_divergence().unwrap()`.
    ///
    /// Each `push()` re-runs `N = ceil(strength * 10)` Gauss-Seidel Laplacian
    /// sweeps from scratch on the raw polyline. Between frames, the only
    /// inputs that differ are the new tip and the now-interior previous tip
    /// (formerly pinned). Both perturbations sit at indices `>= len - 2`.
    ///
    /// Per sweep, the *backward* influence radius of a Gauss-Seidel Laplacian
    /// pass is exactly 1: forward in-sweep updates do not move information
    /// backward, so a change at index `i` can only affect index `i-1` in the
    /// *next* sweep. After `N` sweeps the backward reach is `N`. The earliest
    /// index that can possibly differ between frames is `len - 2 - N`, so the
    /// distance from the tip (`len - 1`) is at most `N + 1`.
    fn max_divergence_window(&self) -> usize {
        if self.strength == 0.0 {
            return 0;
        }
        let iterations = (self.strength * 10.0).ceil() as usize;
        iterations + 1
    }

    fn clear(&mut self) {
        self.raw_points.clear();
        self.stabilized.clear();
        self.prev_positions.clear();
    }
}

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

    fn make_point(x: f32, y: f32) -> PaintInformation {
        PaintInformation {
            pos: [x, y],
            pressure: 0.5,
            ..Default::default()
        }
    }

    fn make_point_with_pressure(x: f32, y: f32, pressure: f32) -> PaintInformation {
        PaintInformation {
            pos: [x, y],
            pressure,
            ..Default::default()
        }
    }

    #[test]
    fn straight_line_stays_straight() {
        let mut stab = LaplacianStabilizer::new(0.8);
        for i in 0..10 {
            stab.push(make_point(i as f32 * 10.0, 0.0));
        }

        // All points should be on y=0 (straight line is already smooth).
        for pt in stab.stabilized() {
            assert!(pt.pos[1].abs() < 1e-3, "y={} should be ~0", pt.pos[1]);
        }
    }

    #[test]
    fn sharp_turn_is_smoothed() {
        let mut stab = LaplacianStabilizer::new(0.8);

        // Straight right, then sharp turn down.
        for i in 0..5 {
            stab.push(make_point(i as f32 * 10.0, 0.0));
        }
        for i in 1..5 {
            stab.push(make_point(40.0, i as f32 * 10.0));
        }

        // The corner point (40, 0) should be pulled inward by smoothing.
        let corner = &stab.stabilized()[4];
        // It should have moved — either x decreased or y increased.
        let moved = corner.pos[0] < 40.0 - 0.1 || corner.pos[1] > 0.1;
        assert!(
            moved,
            "corner at {:?} should be smoothed away from (40, 0)",
            corner.pos
        );
    }

    #[test]
    fn strength_zero_is_pass_through() {
        let mut stab = LaplacianStabilizer::new(0.0);
        let points: Vec<_> = (0..5)
            .map(|i| make_point(i as f32 * 10.0, (i as f32).sin() * 5.0))
            .collect();

        for pt in &points {
            let result = stab.push(*pt);
            assert!(result.divergence_index.is_none());
        }

        // Output should exactly match input.
        for (orig, stab_pt) in points.iter().zip(stab.stabilized()) {
            assert!((orig.pos[0] - stab_pt.pos[0]).abs() < 1e-6);
            assert!((orig.pos[1] - stab_pt.pos[1]).abs() < 1e-6);
        }
    }

    #[test]
    fn first_and_last_pinned() {
        let mut stab = LaplacianStabilizer::new(1.0);

        // Zigzag pattern.
        stab.push(make_point(0.0, 0.0));
        stab.push(make_point(10.0, 20.0));
        stab.push(make_point(20.0, -20.0));
        stab.push(make_point(30.0, 20.0));
        stab.push(make_point(40.0, 0.0));

        let s = stab.stabilized();
        assert!(
            (s[0].pos[0] - 0.0).abs() < 1e-6,
            "first point must be pinned"
        );
        assert!(
            (s[0].pos[1] - 0.0).abs() < 1e-6,
            "first point must be pinned"
        );
        let last = s.last().unwrap();
        assert!(
            (last.pos[0] - 40.0).abs() < 1e-6,
            "last point must be pinned"
        );
        assert!(
            (last.pos[1] - 0.0).abs() < 1e-6,
            "last point must be pinned"
        );
    }

    #[test]
    fn divergence_detected_near_turn() {
        let mut stab = LaplacianStabilizer::new(0.5);

        // Build a straight stroke.
        for i in 0..10 {
            stab.push(make_point(i as f32 * 10.0, 0.0));
        }

        // Add a sharp turn — this should cause divergence near the end, not at the beginning.
        let result = stab.push(make_point(90.0, 30.0));
        if let Some(div) = result.divergence_index {
            assert!(
                div > 2,
                "divergence at {div} should be near the turn, not at the start"
            );
        }
    }

    #[test]
    fn sensor_values_smoothed() {
        let mut stab = LaplacianStabilizer::new(0.8);

        // Pressure spike in the middle.
        stab.push(make_point_with_pressure(0.0, 0.0, 0.3));
        stab.push(make_point_with_pressure(10.0, 0.0, 0.3));
        stab.push(make_point_with_pressure(20.0, 0.0, 1.0)); // spike
        stab.push(make_point_with_pressure(30.0, 0.0, 0.3));
        stab.push(make_point_with_pressure(40.0, 0.0, 0.3));

        let s = stab.stabilized();
        // The spike should be smoothed down.
        assert!(
            s[2].pressure < 0.95,
            "pressure spike at {} should be smoothed",
            s[2].pressure
        );
        // Neighbors should be pulled up slightly.
        assert!(
            s[1].pressure > 0.3,
            "pressure {} should be pulled toward spike",
            s[1].pressure
        );
        assert!(
            s[3].pressure > 0.3,
            "pressure {} should be pulled toward spike",
            s[3].pressure
        );
    }

    #[test]
    fn higher_strength_smooths_more() {
        fn corner_displacement(strength: f32) -> f32 {
            let mut stab = LaplacianStabilizer::new(strength);
            for i in 0..5 {
                stab.push(make_point(i as f32 * 10.0, 0.0));
            }
            for i in 1..5 {
                stab.push(make_point(40.0, i as f32 * 10.0));
            }
            let corner = &stab.stabilized()[4];
            let dx = corner.pos[0] - 40.0;
            let dy = corner.pos[1] - 0.0;
            (dx * dx + dy * dy).sqrt()
        }

        let low = corner_displacement(0.2);
        let high = corner_displacement(0.8);
        assert!(
            high > low,
            "higher strength ({high}) should displace corner more than lower ({low})"
        );
    }

    /// `max_divergence_window` is the contract the checkpoint ring relies on
    /// for coverage. Lock the value to the influence-radius derivation so any
    /// future change is forced through this test (and the documentation).
    #[test]
    fn max_divergence_window_matches_influence_radius() {
        // strength = 0 → pass-through, no relaxation, no divergence window.
        assert_eq!(LaplacianStabilizer::new(0.0).max_divergence_window(), 0);

        // iterations = ceil(strength * 10); window = iterations + 1.
        for (strength, expected_iters) in [(0.1f32, 1u32), (0.25, 3), (0.5, 5), (0.8, 8), (1.0, 10)]
        {
            let stab = LaplacianStabilizer::new(strength);
            assert_eq!(
                stab.max_divergence_window(),
                expected_iters as usize + 1,
                "strength={strength}: window should be iterations+1"
            );
        }
    }

    /// Regression for the prior "find_divergence can return `Some(0)`" bug.
    /// The detector and the bound are co-derived now; this test enforces that
    /// `find_divergence` never returns an index further from the tip than
    /// `max_divergence_window` indices behind it.
    #[test]
    fn find_divergence_respects_max_window() {
        let mut stab = LaplacianStabilizer::new(1.0);
        let max_back = stab.max_divergence_window();

        // Build a long stroke with curvature so divergence is detected on
        // most pushes (every interior point shifts slightly when a new tip
        // is pinned somewhere else).
        for i in 0..200 {
            let t = i as f32;
            let x = t * 4.0;
            let y = (t * 0.15).sin() * 30.0;
            stab.push(make_point(x, y));

            let len = stab.stabilized().len();
            let earliest = len.saturating_sub(max_back + 1);
            if let Some(div) = stab.find_divergence() {
                assert!(
                    div >= earliest,
                    "find_divergence returned {div} at len={len}, earliest allowed = {earliest} \
                     (max_divergence_window = {max_back})"
                );
            }
        }
    }
}