ringgrid 0.5.3

//! Plane-to-image homography estimation via DLT with Hartley normalization.
//!
//! Provides:
//! - Direct Linear Transform (DLT) from ≥4 point correspondences.
//! - RANSAC wrapper for outlier-robust fitting.
//! - Reprojection error computation.

use nalgebra::{Matrix3, Point2, Vector3};
use projective_grid::estimate_homography as pg_estimate_homography;

// ── Error type ───────────────────────────────────────────────────────────

/// Errors produced by homography estimation/refinement routines.
#[derive(Debug, Clone, PartialEq)]
pub enum HomographyError {
    /// Too few correspondences were provided.
    TooFewPoints {
        /// Required minimum number of points.
        needed: usize,
        /// Number of points provided.
        got: usize,
    },
    /// Numerical failure (for example, singular matrices).
    NumericalFailure(String),
    /// RANSAC did not find enough inliers.
    InsufficientInliers {
        /// Required minimum number of inliers.
        needed: usize,
        /// Number of inliers found.
        found: usize,
    },
}

impl std::fmt::Display for HomographyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::TooFewPoints { needed, got } => {
                write!(f, "too few points: need {}, got {}", needed, got)
            }
            Self::NumericalFailure(msg) => write!(f, "numerical failure: {}", msg),
            Self::InsufficientInliers { needed, found } => {
                write!(f, "insufficient inliers: need {}, found {}", needed, found)
            }
        }
    }
}

impl std::error::Error for HomographyError {}

// ── Output statistics ────────────────────────────────────────────────────

/// RANSAC statistics for homography fitting.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct RansacStats {
    /// Number of decoded candidates fed to RANSAC.
    pub n_candidates: usize,
    /// Number of inliers after RANSAC.
    pub n_inliers: usize,
    /// Inlier threshold in working-frame pixels.
    pub threshold_px: f64,
    /// Mean reprojection error of inliers (working-frame pixels).
    pub mean_err_px: f64,
    /// 95th percentile reprojection error (working-frame pixels).
    pub p95_err_px: f64,
}

// ── Projection ───────────────────────────────────────────────────────────

/// Project a 2D point through a 3×3 homography: H * [x, y, 1]^T → [u, v].
pub fn homography_project(h: &Matrix3<f64>, x: f64, y: f64) -> [f64; 2] {
    let p = h * Vector3::new(x, y, 1.0);
    if p[2].abs() < 1e-15 {
        return [f64::NAN, f64::NAN];
    }
    [p[0] / p[2], p[1] / p[2]]
}

/// Reprojection error: ||project(H, src) - dst||.
pub fn homography_reprojection_error(h: &Matrix3<f64>, src: &[f64; 2], dst: &[f64; 2]) -> f64 {
    let p = homography_project(h, src[0], src[1]);
    let dx = p[0] - dst[0];
    let dy = p[1] - dst[1];
    (dx * dx + dy * dy).sqrt()
}

// ── DLT ──────────────────────────────────────────────────────────────────

/// Estimate homography from ≥4 point correspondences using DLT.
///
/// `src`: source points (e.g., board coordinates in mm).
/// `dst`: destination points (e.g., image coordinates in pixels).
///
/// Returns the 3×3 homography H such that dst ≈ project(H, src).
///
/// Delegates to `projective_grid::estimate_homography` (Hartley-normalized
/// DLT with a faster LU-based solver for the exact 4-point case).
pub fn estimate_homography_dlt(
    src: &[[f64; 2]],
    dst: &[[f64; 2]],
) -> Result<Matrix3<f64>, HomographyError> {
    let n = src.len();
    if n < 4 || dst.len() < 4 {
        return Err(HomographyError::TooFewPoints {
            needed: 4,
            got: n.min(dst.len()),
        });
    }
    if src.len() != dst.len() {
        return Err(HomographyError::NumericalFailure(
            "src and dst must have the same length".into(),
        ));
    }

    let src_pts: Vec<Point2<f64>> = src.iter().map(|p| Point2::new(p[0], p[1])).collect();
    let dst_pts: Vec<Point2<f64>> = dst.iter().map(|p| Point2::new(p[0], p[1])).collect();

    pg_estimate_homography(&src_pts, &dst_pts)
        .map(|h| h.h)
        .ok_or_else(|| HomographyError::NumericalFailure("DLT estimation failed".into()))
}

// ── RANSAC ───────────────────────────────────────────────────────────────

/// RANSAC configuration for homography fitting.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(default)]
pub struct RansacHomographyConfig {
    /// Maximum number of RANSAC iterations.
    pub max_iters: usize,
    /// Inlier threshold (reprojection error in pixels).
    pub inlier_threshold: f64,
    /// Minimum number of inliers for a valid model.
    pub min_inliers: usize,
    /// Random seed.
    pub seed: u64,
}

impl Default for RansacHomographyConfig {
    fn default() -> Self {
        Self {
            max_iters: 2000,
            inlier_threshold: 5.0,
            min_inliers: 6,
            seed: 0,
        }
    }
}

/// Result of RANSAC homography fitting.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct RansacHomographyResult {
    /// The fitted homography.
    pub h: Matrix3<f64>,
    /// Boolean mask: true for inliers.
    pub inlier_mask: Vec<bool>,
    /// Number of inliers.
    pub n_inliers: usize,
    /// Per-inlier reprojection errors (only for inliers; others are 0).
    pub errors: Vec<f64>,
}

/// Sample 4 distinct random indices from `[0, n)`.
fn sample_4_distinct(n: usize, rng: &mut rand::rngs::StdRng) -> [usize; 4] {
    use rand::RngExt;
    let mut indices = [0usize; 4];
    for _ in 0..200 {
        for idx in &mut indices {
            *idx = rng.random_range(0..n);
        }
        let mut ok = true;
        for i in 0..4 {
            for j in (i + 1)..4 {
                if indices[i] == indices[j] {
                    ok = false;
                }
            }
        }
        if ok {
            return indices;
        }
    }
    indices
}

/// Fit homography with RANSAC.
///
/// `src`: source points (board coords).
/// `dst`: destination points (image coords).
pub fn fit_homography_ransac(
    src: &[[f64; 2]],
    dst: &[[f64; 2]],
    config: &RansacHomographyConfig,
) -> Result<RansacHomographyResult, HomographyError> {
    let n = src.len();
    if n < 4 {
        return Err(HomographyError::TooFewPoints { needed: 4, got: n });
    }

    use rand::SeedableRng;
    let mut rng = rand::rngs::StdRng::seed_from_u64(config.seed);

    let mut best_inliers = 0usize;
    let mut best_mask: Vec<bool> = vec![false; n];
    let mut best_h = Matrix3::identity();
    let mut mask = vec![false; n];
    let mut adaptive_limit = config.max_iters;

    for iter in 0..config.max_iters {
        if iter >= adaptive_limit {
            break;
        }

        let indices = sample_4_distinct(n, &mut rng);
        let s4: Vec<[f64; 2]> = indices.iter().map(|&i| src[i]).collect();
        let d4: Vec<[f64; 2]> = indices.iter().map(|&i| dst[i]).collect();

        let h = match estimate_homography_dlt(&s4, &d4) {
            Ok(h) => h,
            Err(_) => continue,
        };

        // Count inliers (reuse mask buffer)
        mask.fill(false);
        let mut count = 0usize;
        for i in 0..n {
            let err = homography_reprojection_error(&h, &src[i], &dst[i]);
            if err < config.inlier_threshold {
                mask[i] = true;
                count += 1;
            }
        }

        if count > best_inliers {
            best_inliers = count;
            best_mask.copy_from_slice(&mask);
            best_h = h;

            // Adaptive iteration limit (Hartley & Zisserman Algorithm 4.6)
            // confidence = 99.99%, model DoF = 4 (homography from 4 points)
            let w = count as f64 / n as f64;
            if w > 0.0 {
                let p_fail = (1.0 - w.powi(4)).max(1e-15);
                let needed = (1e-4_f64.ln() / p_fail.ln()).ceil() as usize;
                adaptive_limit = needed.max(50).min(config.max_iters);
            }

            // Early exit if >90% inliers
            if count * 10 > n * 9 {
                break;
            }
        }
    }

    if best_inliers < config.min_inliers {
        return Err(HomographyError::InsufficientInliers {
            needed: config.min_inliers,
            found: best_inliers,
        });
    }

    // Refit using all inliers
    let inlier_src: Vec<[f64; 2]> = (0..n).filter(|&i| best_mask[i]).map(|i| src[i]).collect();
    let inlier_dst: Vec<[f64; 2]> = (0..n).filter(|&i| best_mask[i]).map(|i| dst[i]).collect();

    let h_refit = estimate_homography_dlt(&inlier_src, &inlier_dst).unwrap_or(best_h);

    // Recompute errors and mask with refined H
    let mut final_mask = vec![false; n];
    let mut errors = vec![0.0f64; n];
    let mut final_inliers = 0usize;
    for i in 0..n {
        let err = homography_reprojection_error(&h_refit, &src[i], &dst[i]);
        errors[i] = err;
        if err < config.inlier_threshold {
            final_mask[i] = true;
            final_inliers += 1;
        }
    }

    Ok(RansacHomographyResult {
        h: h_refit,
        inlier_mask: final_mask,
        n_inliers: final_inliers,
        errors,
    })
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use rand::{RngExt, SeedableRng};

    fn make_test_homography() -> Matrix3<f64> {
        // Scale + translate + mild perspective
        Matrix3::new(3.5, 0.1, 640.0, -0.05, 3.3, 480.0, 0.0001, -0.00005, 1.0)
    }

    #[test]
    fn test_dlt_exact_4points() {
        let h_true = make_test_homography();
        let src = [[0.0, 0.0], [100.0, 0.0], [100.0, 100.0], [0.0, 100.0]];
        let dst: Vec<[f64; 2]> = src
            .iter()
            .map(|s| homography_project(&h_true, s[0], s[1]))
            .collect();

        let h_est = estimate_homography_dlt(&src, &dst).unwrap();

        // Check reprojection of all 4 points
        for (s, d) in src.iter().zip(&dst) {
            let err = homography_reprojection_error(&h_est, s, d);
            assert!(err < 1e-6, "reprojection error too large: {}", err);
        }
    }

    #[test]
    fn test_dlt_overdetermined() {
        let h_true = make_test_homography();
        // Grid of 5x5 points
        let mut src = Vec::new();
        let mut dst = Vec::new();
        for i in 0..5 {
            for j in 0..5 {
                let s = [i as f64 * 20.0, j as f64 * 20.0];
                let d = homography_project(&h_true, s[0], s[1]);
                src.push(s);
                dst.push(d);
            }
        }

        let h_est = estimate_homography_dlt(&src, &dst).unwrap();

        for (s, d) in src.iter().zip(&dst) {
            let err = homography_reprojection_error(&h_est, s, d);
            assert!(err < 1e-6, "reprojection error: {}", err);
        }
    }

    #[test]
    fn test_ransac_with_outliers() {
        let h_true = make_test_homography();
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);

        // 20 inlier points
        let mut src = Vec::new();
        let mut dst = Vec::new();
        for i in 0..20 {
            let s = [(i % 5) as f64 * 30.0, (i / 5) as f64 * 30.0];
            let d = homography_project(&h_true, s[0], s[1]);
            // Add small noise
            let d = [
                d[0] + rng.random_range(-0.5..0.5),
                d[1] + rng.random_range(-0.5..0.5),
            ];
            src.push(s);
            dst.push(d);
        }

        // 8 outliers
        for _ in 0..8 {
            let s = [rng.random_range(0.0..100.0), rng.random_range(0.0..100.0)];
            let d = [rng.random_range(0.0..1280.0), rng.random_range(0.0..960.0)];
            src.push(s);
            dst.push(d);
        }

        let config = RansacHomographyConfig {
            max_iters: 2000,
            inlier_threshold: 3.0,
            min_inliers: 6,
            seed: 99,
        };

        let result = fit_homography_ransac(&src, &dst, &config).unwrap();

        // Should find at least 18 of the 20 inliers
        assert!(result.n_inliers >= 18, "only {} inliers", result.n_inliers);

        // Check that reprojection errors for true inliers are small
        for i in 0..20 {
            let err = homography_reprojection_error(&result.h, &src[i], &dst[i]);
            assert!(err < 5.0, "inlier {} has error {}", i, err);
        }
    }

    #[test]
    fn test_project_roundtrip() {
        let h = make_test_homography();
        let h_inv = h.try_inverse().unwrap();

        let p = [50.0, 75.0];
        let q = homography_project(&h, p[0], p[1]);
        let p_back = homography_project(&h_inv, q[0], q[1]);

        assert_relative_eq!(p[0], p_back[0], epsilon = 1e-8);
        assert_relative_eq!(p[1], p_back[1], epsilon = 1e-8);
    }

    #[test]
    fn test_too_few_points() {
        let src = [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]];
        let dst = [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]];
        let result = estimate_homography_dlt(&src, &dst);
        assert!(result.is_err());
    }
}