autopilot/crates/movement_detector/src/internal/ego_motion.rs

//! AZ-662 — Ego-motion estimator + telemetry-skew gate.
//!
//! `EgoMotionEstimator::estimate` checks gimbal/UAV timestamp skew against the
//! per-zoom-band tolerance, then runs OpenCV Lucas–Kanade optical-flow +
//! RANSAC homography on consecutive grayscale frames to recover camera motion.

use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc,
};

use opencv::{core::Mat, prelude::*};

use shared::models::{
    frame::Frame,
    gimbal::GimbalState,
    movement::ZoomBand,
    telemetry::UavTelemetry,
};

use super::{
    optical_flow::{self, FlowError},
    telemetry_sync::{self, SkewExceeded},
    zoom_bands::zoom_band_from_level,
};

/// Per-frame ego-motion recovered from optical flow.
#[derive(Debug, Clone)]
pub struct EgoMotion {
    /// Row-major 3×3 homography mapping the previous frame's coordinates to
    /// the current frame's coordinates (camera ego-motion).
    pub homography: [f64; 9],
    /// Mean reprojection residual across inlier feature tracks (pixels).
    pub residual_motion_magnitude: f32,
    pub zoom_band: ZoomBand,
}

/// Error variants returned by `EgoMotionEstimator::estimate`.
#[derive(Debug)]
pub enum EgoMotionError {
    /// Frame ↔ gimbal or frame ↔ UAV timestamp skew exceeded the per-band
    /// tolerance. The affected frame must not be used for ego-motion.
    SkewExceeded(SkewExceeded),
    /// The current frame is degenerate (saturated, blank, or featureless).
    /// The frame is stored internally so the next call can resume.
    OpticalFlowDegenerate,
    /// No previous frame has been received yet; the current frame is stored
    /// as the reference for the next call.
    NoPreviousFrame,
    Internal(String),
}

impl From<SkewExceeded> for EgoMotionError {
    fn from(e: SkewExceeded) -> Self {
        EgoMotionError::SkewExceeded(e)
    }
}

/// Atomic health counters exposed through `MovementDetectorHandle::health()`.
pub struct EgoMotionCounters {
    pub telemetry_skew_drops_zoomed_out: AtomicU64,
    pub telemetry_skew_drops_zoomed_in: AtomicU64,
    pub optical_flow_degenerate_total: AtomicU64,
}

impl EgoMotionCounters {
    pub fn new() -> Self {
        Self {
            telemetry_skew_drops_zoomed_out: AtomicU64::new(0),
            telemetry_skew_drops_zoomed_in: AtomicU64::new(0),
            optical_flow_degenerate_total: AtomicU64::new(0),
        }
    }

    pub fn skew_drops(&self, band: ZoomBand) -> u64 {
        match band {
            ZoomBand::ZoomedOut => {
                self.telemetry_skew_drops_zoomed_out.load(Ordering::Relaxed)
            }
            ZoomBand::ZoomedIn => {
                self.telemetry_skew_drops_zoomed_in.load(Ordering::Relaxed)
            }
        }
    }

    pub fn skew_drops_total(&self) -> u64 {
        self.skew_drops(ZoomBand::ZoomedOut) + self.skew_drops(ZoomBand::ZoomedIn)
    }

    pub fn degenerate_total(&self) -> u64 {
        self.optical_flow_degenerate_total.load(Ordering::Relaxed)
    }

    fn inc_skew_drop(&self, band: ZoomBand) {
        match band {
            ZoomBand::ZoomedOut => {
                self.telemetry_skew_drops_zoomed_out.fetch_add(1, Ordering::Relaxed);
            }
            ZoomBand::ZoomedIn => {
                self.telemetry_skew_drops_zoomed_in.fetch_add(1, Ordering::Relaxed);
            }
        }
    }

    fn inc_degenerate(&self) {
        self.optical_flow_degenerate_total.fetch_add(1, Ordering::Relaxed);
    }
}

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

/// Stateful per-frame ego-motion estimator.
///
/// Call `estimate` once per frame in arrival order. The estimator keeps the
/// previous frame's grayscale Mat internally; the first call always returns
/// `Err(NoPreviousFrame)` and stores the frame as the reference.
pub struct EgoMotionEstimator {
    prev_gray: Option<Mat>,
    counters: Arc<EgoMotionCounters>,
}

impl EgoMotionEstimator {
    pub fn new(counters: Arc<EgoMotionCounters>) -> Self {
        Self { prev_gray: None, counters }
    }

    pub fn counters(&self) -> &Arc<EgoMotionCounters> {
        &self.counters
    }

    /// Estimate ego-motion for `frame` relative to the previous accepted frame.
    ///
    /// Processing order:
    /// 1. Telemetry-skew gate (increments `telemetry_skew_drops_total` on miss).
    /// 2. Convert to grayscale.
    /// 3. Degenerate-frame detection (increments `optical_flow_degenerate_total`).
    /// 4. Require a previous accepted frame; store current if none.
    /// 5. LK optical flow + RANSAC homography.
    pub fn estimate(
        &mut self,
        frame: &Frame,
        gimbal_state: &GimbalState,
        uav_telemetry: &UavTelemetry,
    ) -> Result<EgoMotion, EgoMotionError> {
        let zoom_band = zoom_band_from_level(gimbal_state.zoom);

        // 1. Skew gate.
        telemetry_sync::check_skew(
            frame.capture_ts_monotonic_ns,
            gimbal_state.ts_monotonic_ns,
            uav_telemetry.monotonic_ts_ns,
            zoom_band,
        )
        .map_err(|e| {
            self.counters.inc_skew_drop(zoom_band);
            EgoMotionError::SkewExceeded(e)
        })?;

        // 2. Grayscale conversion.
        let curr_gray = optical_flow::frame_to_gray(frame)
            .map_err(|e| EgoMotionError::Internal(e.message))?;

        // 3. Degenerate check — runs before the prev-frame guard so a
        //    saturated frame still stores itself and returns a clear error.
        if optical_flow::is_degenerate(&curr_gray)
            .unwrap_or(false)
        {
            self.counters.inc_degenerate();
            self.prev_gray = Some(curr_gray);
            return Err(EgoMotionError::OpticalFlowDegenerate);
        }

        // 4. Need a previous frame for optical flow.
        let prev_gray = match self.prev_gray.take() {
            None => {
                self.prev_gray = Some(curr_gray);
                return Err(EgoMotionError::NoPreviousFrame);
            }
            Some(p) => p,
        };

        // 5. Optical flow → homography.
        let result = optical_flow::estimate_homography(&prev_gray, &curr_gray);
        self.prev_gray = Some(curr_gray);

        match result {
            Ok(hr) => Ok(EgoMotion {
                homography: hr.h,
                residual_motion_magnitude: hr.residual_magnitude_px,
                zoom_band,
            }),
            Err(FlowError::Degenerate | FlowError::InsufficientFeatures) => {
                self.counters.inc_degenerate();
                Err(EgoMotionError::OpticalFlowDegenerate)
            }
            Err(FlowError::Internal(msg)) => Err(EgoMotionError::Internal(msg)),
        }
    }
}

// ── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use bytes::Bytes;
    use opencv::{
        core::{Mat, Scalar, CV_8UC1},
        prelude::*,
    };

    use shared::models::{
        frame::{Frame, PixelFormat},
        gimbal::GimbalState,
        movement::ZoomBand,
        telemetry::UavTelemetry,
    };

    use super::*;

    // ── helpers ──────────────────────────────────────────────────────────────

    /// Build a 1-channel Mat filled by `fill(row, col)`.
    fn make_gray_mat(
        size: i32,
        fill: impl Fn(i32, i32) -> u8,
    ) -> opencv::Result<Mat> {
        let mut mat =
            Mat::new_rows_cols_with_default(size, size, CV_8UC1, Scalar::all(0.0))?;
        for r in 0..size {
            for c in 0..size {
                *mat.at_2d_mut::<u8>(r, c)? = fill(r, c);
            }
        }
        Ok(mat)
    }

    /// Checkerboard with 8-pixel blocks, optionally shifted right by `offset_x`.
    fn checkerboard(size: i32, offset_x: i32) -> opencv::Result<Mat> {
        make_gray_mat(size, |r, c| {
            let sc = c - offset_x;
            if sc < 0 || sc >= size {
                128
            } else if (sc / 8 + r / 8) % 2 == 0 {
                200
            } else {
                50
            }
        })
    }

    /// Wrap a 1-channel Mat as a Nv12 `Frame` (Y-plane only — sufficient for
    /// `frame_to_gray` which reads only the first w×h bytes).
    fn mat_to_frame(mat: &Mat, ts_ns: u64) -> opencv::Result<Frame> {
        let h = mat.rows() as u32;
        let w = mat.cols() as u32;
        let total = (w * h) as usize;
        let mut pixels = vec![0u8; total];
        for r in 0..h as i32 {
            for c in 0..w as i32 {
                pixels[(r * w as i32 + c) as usize] = *mat.at_2d::<u8>(r, c)?;
            }
        }
        Ok(Frame {
            seq: 0,
            capture_ts_monotonic_ns: ts_ns,
            decode_ts_monotonic_ns: ts_ns,
            pixels: Arc::new(Bytes::from(pixels)),
            width: w,
            height: h,
            pix_fmt: PixelFormat::Nv12,
            ai_locked: false,
        })
    }

    fn synced_gimbal(ts_ns: u64) -> GimbalState {
        GimbalState {
            yaw: 0.0,
            pitch: -30.0,
            zoom: 1.0,
            ts_monotonic_ns: ts_ns,
            command_in_flight: false,
        }
    }

    fn synced_uav(ts_ns: u64) -> UavTelemetry {
        UavTelemetry { monotonic_ts_ns: ts_ns, ..UavTelemetry::empty() }
    }

    // ── AC-1: synthetic pure pan — residual ≈ 0 ──────────────────────────────

    #[test]
    fn ac1_pure_pan_residual_near_zero() -> opencv::Result<()> {
        let counters = Arc::new(EgoMotionCounters::new());
        let mut est = EgoMotionEstimator::new(Arc::clone(&counters));

        let size = 200;
        let dx = 8i32; // one checkerboard block = well-defined shift

        let mat1 = checkerboard(size, 0)?;
        let mat2 = checkerboard(size, dx)?;

        let t0 = 1_000_000_000u64;
        let frame1 = mat_to_frame(&mat1, t0)?;
        let frame2 = mat_to_frame(&mat2, t0 + 33_000_000)?; // +33 ms (30 fps)

        let gimbal = synced_gimbal(t0);
        let uav = synced_uav(t0);

        // First call stores prev; NoPreviousFrame is expected.
        assert!(matches!(est.estimate(&frame1, &gimbal, &uav), Err(EgoMotionError::NoPreviousFrame)));

        let gimbal2 = synced_gimbal(t0 + 33_000_000);
        let uav2 = synced_uav(t0 + 33_000_000);
        let ego = est.estimate(&frame2, &gimbal2, &uav2)
            .expect("estimate should succeed on second call");

        // X-translation H[0][2] should approximate dx within ±2 px.
        let h02 = ego.homography[2];
        assert!(
            h02.abs() > 0.5 && (h02 - dx as f64).abs() < 2.5,
            "H[0][2] = {h02:.2}, expected ≈ {dx}"
        );
        // Residual should be near zero for a pure rigid shift.
        assert!(
            ego.residual_motion_magnitude < 3.0,
            "residual = {:.2} px, expected < 3.0",
            ego.residual_motion_magnitude
        );
        assert_eq!(ego.zoom_band, ZoomBand::ZoomedOut);
        Ok(())
    }

    // ── AC-2: telemetry skew above zoom-out tolerance → SkewExceeded ─────────

    #[test]
    fn ac2_skew_above_zoom_out_tolerance_dropped() -> opencv::Result<()> {
        let counters = Arc::new(EgoMotionCounters::new());
        let mut est = EgoMotionEstimator::new(Arc::clone(&counters));

        let frame_ts = 1_000_000_000u64;
        let frame = mat_to_frame(&checkerboard(100, 0)?, frame_ts)?;

        // Gimbal timestamp 200 ms ahead of frame; tolerance = 50 ms.
        let gimbal = GimbalState {
            zoom: 1.0, // zoomed_out → 50 ms tolerance
            ts_monotonic_ns: frame_ts + 200_000_000,
            yaw: 0.0,
            pitch: -30.0,
            command_in_flight: false,
        };
        let uav = synced_uav(frame_ts);

        assert!(matches!(
            est.estimate(&frame, &gimbal, &uav),
            Err(EgoMotionError::SkewExceeded(_))
        ));
        assert_eq!(counters.skew_drops(ZoomBand::ZoomedOut), 1);
        Ok(())
    }

    // ── AC-3: fully-saturated white frame → OpticalFlowDegenerate ────────────

    #[test]
    fn ac3_degenerate_white_frame() -> opencv::Result<()> {
        let counters = Arc::new(EgoMotionCounters::new());
        let mut est = EgoMotionEstimator::new(Arc::clone(&counters));

        let ts = 1_000_000_000u64;
        let white_mat =
            Mat::new_rows_cols_with_default(100, 100, CV_8UC1, Scalar::all(255.0))?;
        let frame = mat_to_frame(&white_mat, ts)?;
        let gimbal = synced_gimbal(ts);
        let uav = synced_uav(ts);

        assert!(matches!(
            est.estimate(&frame, &gimbal, &uav),
            Err(EgoMotionError::OpticalFlowDegenerate)
        ));
        assert_eq!(counters.degenerate_total(), 1);
        Ok(())
    }
}