kornia-3d 0.1.14

//! RANSAC-based robust wrapper for PnP solvers.

use super::ops::{intrinsics_as_vectors, project_sq_error};
use super::{solve_pnp, PnPMethod};
use super::{PnPError, PnPResult};
use kornia_algebra::{Mat3AF32, Vec2F32, Vec3AF32};
use kornia_imgproc::calibration::distortion::PolynomialDistortion;
use rand::seq::SliceRandom;
use rand::{rngs::StdRng, SeedableRng};
use thiserror::Error;

const MIN_CORRESPONDENCES: usize = 4; // Minimum 2D-3D pairs required by PnP
const EPNP_MIN_SAMPLE_SIZE: usize = 5; // Minimal sample size for EPnP (unless only 4 points available)

const DEFAULT_MAX_ITERATIONS: usize = 100;
const DEFAULT_REPROJ_THRESHOLD_PX: f32 = 8.0;
const DEFAULT_CONFIDENCE: f32 = 0.99;

const EPS_PROB_MIN: f32 = 1e-6; // Guard for tiny probabilities
const EPS_LOG_GUARD: f32 = 1e-12; // Guard to avoid log(0) and log(1)
const HIGH_INLIER_RATIO_STOP: f32 = 0.95; // Early stop when inlier ratio is very high

/// Error type for the RANSAC-based PnP solver.
///
/// This encapsulates both underlying base-solver errors (`PnPError`) as well as
/// RANSAC-specific failure modes.
#[derive(Debug, Error)]
pub enum PnPRansacError {
    /// Error returned by the underlying PnP solver used inside RANSAC.
    #[error(transparent)]
    Base(#[from] PnPError),

    /// RANSAC failed to find a valid model with enough inliers.
    #[error("RANSAC found insufficient inliers: required {required}, got {actual}")]
    InsufficientInliers {
        /// Minimum number of inliers required to accept a model
        required: usize,
        /// Actual number of inliers found
        actual: usize,
    },
}

/// Parameters for RANSAC over PnP.
#[derive(Debug, Clone)]
pub struct RansacParams {
    /// Maximum number of RANSAC iterations.
    pub max_iterations: usize,
    /// Pixel error threshold to classify an observation as an inlier.
    pub reproj_threshold_px: f32,
    /// Desired probability that at least one sample set is outlier-free.
    pub confidence: f32,
    /// Optional fixed seed for reproducible sampling.
    pub random_seed: Option<u64>,
    /// Whether to refit on all inliers using the base solver.
    pub refine: bool,
}

impl Default for RansacParams {
    fn default() -> Self {
        Self {
            max_iterations: DEFAULT_MAX_ITERATIONS,
            reproj_threshold_px: DEFAULT_REPROJ_THRESHOLD_PX,
            confidence: DEFAULT_CONFIDENCE,
            random_seed: None,
            refine: true,
        }
    }
}

/// RANSAC result for PnP.
#[derive(Debug, Clone)]
pub struct PnPRansacResult {
    /// Best pose found by RANSAC.
    pub pose: PnPResult,
    /// Indices of inlier correspondences.
    pub inliers: Vec<usize>,
}

/// Solve PnP robustly using a legacy RANSAC loop around a base PnP method (e.g., EPnP).
///
/// - Minimal sample size is 5 for EPnP (4 when only 4 points available).
/// - Scoring uses Euclidean pixel reprojection error.
/// - Iterations adapt from current inlier ratio and desired confidence.
pub fn solve_pnp_ransac(
    world: &[Vec3AF32],
    image: &[Vec2F32],
    k: &Mat3AF32,
    distortion: Option<&PolynomialDistortion>,
    base: PnPMethod,
    params: &RansacParams,
) -> Result<PnPRansacResult, PnPRansacError> {
    let n = world.len();
    if n != image.len() {
        return Err(PnPError::MismatchedArrayLengths {
            left_name: "world points",
            left_len: world.len(),
            right_name: "image points",
            right_len: image.len(),
        }
        .into());
    }
    if n < MIN_CORRESPONDENCES {
        return Err(PnPError::InsufficientCorrespondences {
            required: MIN_CORRESPONDENCES,
            actual: n,
        }
        .into());
    }

    // Minimal set size: EPnP uses 5 points (unless only 4 points available)
    let sample_size: usize = if n == MIN_CORRESPONDENCES {
        MIN_CORRESPONDENCES
    } else {
        EPNP_MIN_SAMPLE_SIZE
    };

    // Precompute intrinsics vectors
    let (intr_x, intr_y) = intrinsics_as_vectors(k);

    // RNG setup
    let mut rng: StdRng = match params.random_seed {
        Some(seed) => StdRng::seed_from_u64(seed),
        None => {
            let mut trng = rand::rng();
            StdRng::from_rng(&mut trng)
        }
    };

    // Working buffers
    let mut indices: Vec<usize> = (0..n).collect();
    let mut best_inliers: Vec<usize> = Vec::new();
    let mut best_pose: Option<PnPResult> = None;
    let mut w_min: Vec<Vec3AF32> = Vec::with_capacity(sample_size);
    let mut i_min: Vec<Vec2F32> = Vec::with_capacity(sample_size);

    let mut iter: usize = 0;
    let mut required_iters = params.max_iterations;

    while iter < required_iters && iter < params.max_iterations {
        iter += 1;

        // Debug: prevent infinite loops
        if iter > params.max_iterations {
            log::warn!("RANSAC: Emergency break after {iter} iterations");
            break;
        }

        // Sample k unique indices without replacement.
        indices.shuffle(&mut rng);
        let sample = &indices[..sample_size];

        // Build minimal subsets
        w_min.clear();
        i_min.clear();
        for &idx in sample.iter() {
            w_min.push(world[idx]);
            i_min.push(image[idx]);
        }

        // Estimate pose on minimal set
        let pose_maybe = solve_pnp(&w_min, &i_min, k, distortion, base.clone());
        let pose_min = match pose_maybe {
            Ok(p) => p,
            Err(_e) => {
                log::debug!("EPnP failed on minimal set");
                continue;
            }
        };

        // Optional cheirality check on minimal set (all positive depths)
        if !sample_all_positive_depths(&pose_min.rotation, &pose_min.translation, &w_min) {
            log::debug!("Cheirality check failed on iteration {iter}");
            continue;
        }

        // Score model on all points
        let (inliers, _total_squared_error) = classify_points(
            world,
            image,
            None,
            None,
            ClassificationParams {
                rotation_matrix: &pose_min.rotation,
                translation_vector: &pose_min.translation,
                camera_intrinsics_x: &intr_x,
                camera_intrinsics_y: &intr_y,
                threshold: Some(params.reproj_threshold_px),
            },
        );

        if inliers.len() > best_inliers.len() {
            best_inliers = inliers;
            best_pose = Some(pose_min);

            // Update required iterations based on current inlier ratio and sample size
            if best_inliers.len() >= sample_size {
                let w = best_inliers.len() as f32 / n as f32;
                let s = sample_size as f32;

                // Avoid numerical issues with very small w
                if w > EPS_PROB_MIN && w < 1.0 {
                    let ws = w.powf(s);
                    if ws < 1.0 - EPS_LOG_GUARD && ws > EPS_LOG_GUARD {
                        // Avoid log(0) and log(1)
                        let log_conf = (1.0 - params.confidence).max(EPS_LOG_GUARD).ln();
                        let log_denom = (1.0 - ws).ln();
                        if log_denom.is_finite() && log_denom.abs() > EPS_LOG_GUARD {
                            let est = (log_conf / log_denom).ceil();

                            if est.is_finite() && est > 0.0 {
                                let est_usize = est.min(params.max_iterations as f32) as usize;
                                if est_usize < required_iters {
                                    required_iters = est_usize;
                                }
                            }
                        }
                    } else if w >= HIGH_INLIER_RATIO_STOP {
                        // Very high inlier ratio (≥95%), we can stop early
                        required_iters = iter;
                    }
                }
            }
        }
    }

    // Validate and optionally refine
    if best_inliers.len() < MIN_CORRESPONDENCES {
        let err = PnPRansacError::InsufficientInliers {
            required: MIN_CORRESPONDENCES,
            actual: best_inliers.len(),
        };
        return Err(err);
    }

    let mut final_pose = if params.refine {
        // Refit on all inliers using the base solver.
        let mut w_all = Vec::with_capacity(best_inliers.len());
        let mut i_all = Vec::with_capacity(best_inliers.len());
        for &idx in &best_inliers {
            w_all.push(world[idx]);
            i_all.push(image[idx]);
        }
        solve_pnp(&w_all, &i_all, k, distortion, base.clone())?
    } else {
        match best_pose {
            Some(p) => p,
            None => {
                return Err(PnPError::SvdFailed(
                    "RANSAC failed to produce a pose despite sufficient inliers".to_string(),
                )
                .into());
            }
        }
    };

    // Recompute reprojection error on inliers only
    let (_inliers, sum_sq_inliers) = classify_points(
        world,
        image,
        None,
        Some(&best_inliers),
        ClassificationParams {
            rotation_matrix: &final_pose.rotation,
            translation_vector: &final_pose.translation,
            camera_intrinsics_x: &intr_x,
            camera_intrinsics_y: &intr_y,
            threshold: None,
        },
    );
    let rmse = if !best_inliers.is_empty() {
        (sum_sq_inliers / best_inliers.len() as f32).sqrt()
    } else {
        0.0
    };
    final_pose.reproj_rmse = Some(rmse);

    Ok(PnPRansacResult {
        pose: final_pose,
        inliers: best_inliers,
    })
}

fn sample_all_positive_depths(r: &Mat3AF32, t: &Vec3AF32, world: &[Vec3AF32]) -> bool {
    world.iter().all(|&pw| {
        let pc = *r * pw + *t;
        pc.z > 0.0
    })
}

/// This function handles both:
/// - Scoring all points against a candidate pose (during RANSAC)
/// - Computing final RMSE on a subset of inlier points
struct ClassificationParams<'a> {
    rotation_matrix: &'a Mat3AF32,
    translation_vector: &'a Vec3AF32,
    camera_intrinsics_x: &'a Vec3AF32,
    camera_intrinsics_y: &'a Vec3AF32,
    threshold: Option<f32>,
}

fn classify_points(
    world: &[Vec3AF32],
    image: &[Vec2F32],
    _distortion: Option<&PolynomialDistortion>,
    indices: Option<&[usize]>,
    params: ClassificationParams,
) -> (Vec<usize>, f32) {
    let rotation = params.rotation_matrix;
    let translation = params.translation_vector;

    let mut inliers: Vec<usize> = Vec::new();
    let mut total_squared_error: f32 = 0.0;

    match indices {
        Some(indices) => {
            for &idx in indices {
                if idx >= world.len() || idx >= image.len() {
                    continue;
                }

                if let Some(squared_error) = project_sq_error(
                    &world[idx],
                    &image[idx],
                    rotation,
                    translation,
                    params.camera_intrinsics_x,
                    params.camera_intrinsics_y,
                    true,
                ) {
                    total_squared_error += squared_error;

                    let is_inlier = match params.threshold {
                        Some(thresh) => squared_error.sqrt() < thresh,
                        None => true,
                    };

                    if is_inlier {
                        inliers.push(idx);
                    }
                }
            }
        }
        None => {
            for (idx, (world_point, image_point)) in world.iter().zip(image.iter()).enumerate() {
                if let Some(squared_error) = project_sq_error(
                    world_point,
                    image_point,
                    rotation,
                    translation,
                    params.camera_intrinsics_x,
                    params.camera_intrinsics_y,
                    true,
                ) {
                    total_squared_error += squared_error;

                    let is_inlier = match params.threshold {
                        Some(thresh) => squared_error.sqrt() < thresh,
                        None => true,
                    };

                    if is_inlier {
                        inliers.push(idx);
                    }
                }
            }
        }
    }

    (inliers, total_squared_error)
}

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

    fn k_default() -> Mat3AF32 {
        // Column-major: [fx, 0, 0,  0, fy, 0,  cx, cy, 1]
        Mat3AF32::from_cols_array(&[800.0, 0.0, 0.0, 0.0, 800.0, 0.0, 640.0, 480.0, 1.0])
    }

    #[test]
    fn test_ransac_basic_outliers() -> Result<(), PnPRansacError> {
        // Use the same 6 inlier correspondences from epnp test
        let points_world: [Vec3AF32; 6] = [
            Vec3AF32::new(0.0315, 0.03333, -0.10409),
            Vec3AF32::new(-0.0315, 0.03333, -0.10409),
            Vec3AF32::new(0.0, -0.00102, -0.12977),
            Vec3AF32::new(0.02646, -0.03167, -0.1053),
            Vec3AF32::new(-0.02646, -0.031667, -0.1053),
            Vec3AF32::new(0.0, 0.04515, -0.11033),
        ];
        let mut points_image: Vec<Vec2F32> = vec![
            Vec2F32::new(722.96466, 502.0828),
            Vec2F32::new(669.88837, 498.61877),
            Vec2F32::new(707.0025, 478.48975),
            Vec2F32::new(728.05634, 447.56918),
            Vec2F32::new(682.6069, 443.91776),
            Vec2F32::new(696.4414, 511.96442),
        ];
        // Inject 4 strong outliers
        let mut world = points_world.to_vec();
        for (j, &point_world) in points_world.iter().enumerate().take(4) {
            world.push(point_world);
            points_image.push(Vec2F32::new(
                1200.0 + j as f32 * 5.0,
                -300.0 - j as f32 * 3.0,
            ));
        }
        let k = k_default();

        let params = RansacParams {
            max_iterations: 10,          // Reduce for testing
            reproj_threshold_px: 1000.0, // High threshold needed for this test data with extreme outliers
            confidence: 0.99,
            random_seed: Some(42),
            refine: false,
        };

        let base = PnPMethod::EPnP(EPnPParams::default());

        let res = solve_pnp_ransac(&world, &points_image, &k, None, base, &params)?;
        assert!(res.inliers.len() >= 6); // Should find at least the 6 original inliers
        assert!(res.pose.reproj_rmse.is_some());

        // With extreme outliers, RMSE will be higher, but RANSAC should still work
        let rmse = res.pose.reproj_rmse.unwrap();
        assert!(rmse < 2000.0); // Allow reasonable tolerance for this challenging test data
        Ok(())
    }

    #[test]
    fn test_ransac_perfect_data() -> Result<(), PnPRansacError> {
        // Test with perfect data (no outliers)
        let points_world: [Vec3AF32; 6] = [
            Vec3AF32::new(0.0315, 0.03333, -0.10409),
            Vec3AF32::new(-0.0315, 0.03333, -0.10409),
            Vec3AF32::new(0.0, -0.00102, -0.12977),
            Vec3AF32::new(0.02646, -0.03167, -0.1053),
            Vec3AF32::new(-0.02646, -0.031667, -0.1053),
            Vec3AF32::new(0.0, 0.04515, -0.11033),
        ];
        let points_image: [Vec2F32; 6] = [
            Vec2F32::new(722.96466, 502.0828),
            Vec2F32::new(669.88837, 498.61877),
            Vec2F32::new(707.0025, 478.48975),
            Vec2F32::new(728.05634, 447.56918),
            Vec2F32::new(682.6069, 443.91776),
            Vec2F32::new(696.4414, 511.96442),
        ];
        let k = k_default();

        let params = RansacParams {
            max_iterations: 10,
            reproj_threshold_px: 8.0,
            confidence: 0.99,
            random_seed: Some(42),
            refine: true,
        };

        let base = PnPMethod::EPnP(EPnPParams::default());
        let res = solve_pnp_ransac(&points_world, &points_image, &k, None, base, &params)?;
        assert_eq!(res.inliers.len(), 6); // All points should be inliers
        assert!(res.pose.reproj_rmse.is_some());

        // The test data from epnp tests has some inherent reprojection error
        // (~12 pixels RMSE) which is reasonable given the 8px threshold
        let rmse = res.pose.reproj_rmse.unwrap();
        assert!(rmse < 20.0); // Allow reasonable tolerance for this test data
        Ok(())
    }

    #[test]
    fn test_ransac_minimum_points() -> Result<(), PnPRansacError> {
        // Test with exactly 4 points
        let points_world: [Vec3AF32; 4] = [
            Vec3AF32::new(0.0315, 0.03333, -0.10409),
            Vec3AF32::new(-0.0315, 0.03333, -0.10409),
            Vec3AF32::new(0.0, -0.00102, -0.12977),
            Vec3AF32::new(0.02646, -0.03167, -0.1053),
        ];
        let points_image: [Vec2F32; 4] = [
            Vec2F32::new(722.96466, 502.0828),
            Vec2F32::new(669.88837, 498.61877),
            Vec2F32::new(707.0025, 478.48975),
            Vec2F32::new(728.05634, 447.56918),
        ];
        let k = k_default();

        let params = RansacParams {
            max_iterations: 5,
            reproj_threshold_px: 8.0,
            confidence: 0.99,
            random_seed: Some(42),
            refine: true,
        };

        let base = PnPMethod::EPnP(EPnPParams::default());
        let res = solve_pnp_ransac(&points_world, &points_image, &k, None, base, &params)?;
        assert!(res.inliers.len() >= 4);
        Ok(())
    }

    #[test]
    fn test_ransac_error_cases() {
        let points_world: [Vec3AF32; 3] = [
            Vec3AF32::new(0.0, 0.0, 1.0),
            Vec3AF32::new(1.0, 0.0, 1.0),
            Vec3AF32::new(0.0, 1.0, 1.0),
        ];
        let points_image: [Vec2F32; 3] = [
            Vec2F32::new(100.0, 100.0),
            Vec2F32::new(200.0, 100.0),
            Vec2F32::new(100.0, 200.0),
        ];
        let k = Mat3AF32::from_cols(
            Vec3AF32::new(800.0, 0.0, 0.0),
            Vec3AF32::new(0.0, 800.0, 0.0),
            Vec3AF32::new(400.0, 300.0, 1.0),
        );

        let params = RansacParams::default();
        let base = PnPMethod::EPnP(EPnPParams::default());

        // Should fail with insufficient correspondences
        let result = solve_pnp_ransac(&points_world, &points_image, &k, None, base, &params);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            PnPRansacError::Base(PnPError::InsufficientCorrespondences { .. })
        ));
    }
}