use nalgebra::{DMatrix, Matrix3, RowDVector, SVD, Vector3};

pub const DEFAULT_MAX_ITERATIONS: usize = 1000;
pub const DEFAULT_INLIER_THRESHOLD: f64 = 4.0;
pub const DEFAULT_CONFIDENCE: f64 = 0.9999;
/// The normalized 8-point algorithm needs at least this many correspondences to fit
/// a fundamental matrix at all.
pub const MIN_CORRESPONDENCES: usize = 8;

pub struct RansacResult {
  /// Row-major 3x3 fundamental matrix.
  pub f_matrix: [f64; 9],
  /// Indices (into the input correspondence slices) of inlier matches.
  pub inliers: Vec<usize>,
}

fn normalization_transform(pts: &[(f64, f64)]) -> Matrix3<f64> {
  let n = pts.len() as f64;
  let (sum_x, sum_y) = pts.iter().fold((0.0, 0.0), |(ax, ay), &(x, y)| (ax + x, ay + y));
  let cx = sum_x / n;
  let cy = sum_y / n;
  let mean_dist = pts
    .iter()
    .map(|&(x, y)| ((x - cx).powi(2) + (y - cy).powi(2)).sqrt())
    .sum::<f64>()
    / n;
  let s = if mean_dist > 1e-12 {
    std::f64::consts::SQRT_2 / mean_dist
  } else {
    1.0
  };
  Matrix3::new(s, 0.0, -s * cx, 0.0, s, -s * cy, 0.0, 0.0, 1.0)
}

fn apply_transform(t: &Matrix3<f64>, p: (f64, f64)) -> (f64, f64) {
  let v = t * Vector3::new(p.0, p.1, 1.0);
  (v.x / v.z, v.y / v.z)
}

/// Normalized 8-point algorithm. Works with >=8 correspondences via least-squares.
fn fit_fundamental(pts0: &[(f64, f64)], pts1: &[(f64, f64)]) -> Option<Matrix3<f64>> {
  let n = pts0.len();
  if n < MIN_CORRESPONDENCES || pts1.len() != n {
    return None;
  }

  let t0 = normalization_transform(pts0);
  let t1 = normalization_transform(pts1);
  let n0: Vec<(f64, f64)> = pts0.iter().map(|&p| apply_transform(&t0, p)).collect();
  let n1: Vec<(f64, f64)> = pts1.iter().map(|&p| apply_transform(&t1, p)).collect();

  let mut a = DMatrix::<f64>::zeros(n, 9);
  for i in 0..n {
    let (x0, y0) = n0[i];
    let (x1, y1) = n1[i];
    a.set_row(
      i,
      &RowDVector::from_vec(vec![x1 * x0, x1 * y0, x1, y1 * x0, y1 * y0, y1, x0, y0, 1.0]),
    );
  }

  let svd = SVD::new(a, false, true);
  let v_t = svd.v_t?;
  let idx = argmin(svd.singular_values.as_slice())?;
  let f_row = v_t.row(idx);
  let f_normalized = Matrix3::new(
    f_row[0], f_row[1], f_row[2], f_row[3], f_row[4], f_row[5], f_row[6], f_row[7], f_row[8],
  );

  // Enforce rank-2 by zeroing the smallest singular value of F.
  let svd_f = SVD::new(f_normalized, true, true);
  let u = svd_f.u?;
  let v_t2 = svd_f.v_t?;
  let mut s = svd_f.singular_values;
  let min_idx = argmin(s.as_slice())?;
  s[min_idx] = 0.0;
  let f_rank2 = u * Matrix3::from_diagonal(&s) * v_t2;

  Some(t1.transpose() * f_rank2 * t0)
}

fn argmin(values: &[f64]) -> Option<usize> {
  values
    .iter()
    .enumerate()
    .min_by(|a, b| a.1.partial_cmp(b.1).unwrap())
    .map(|(i, _)| i)
}

fn sampson_distance(f: &Matrix3<f64>, p0: (f64, f64), p1: (f64, f64)) -> f64 {
  let x0 = Vector3::new(p0.0, p0.1, 1.0);
  let x1 = Vector3::new(p1.0, p1.1, 1.0);
  let fx0 = f * x0;
  let ftx1 = f.transpose() * x1;
  let numerator = x1.dot(&fx0).powi(2);
  let denom = fx0.x.powi(2) + fx0.y.powi(2) + ftx1.x.powi(2) + ftx1.y.powi(2);
  if denom < 1e-12 {
    f64::INFINITY
  } else {
    numerator / denom
  }
}

/// RANSAC estimation of the fundamental matrix relating `pts0` to `pts1` (index i in
/// `pts0` corresponds to index i in `pts1`). Returns `None` if fewer than 8
/// correspondences are given or no valid model is found.
pub fn ransac_fundamental_matrix(
  pts0: &[(f32, f32)],
  pts1: &[(f32, f32)],
  max_iterations: usize,
  inlier_threshold: f64,
  confidence: f64,
) -> Option<RansacResult> {
  let n = pts0.len();
  if n < MIN_CORRESPONDENCES || pts1.len() != n {
    return None;
  }

  let pts0f: Vec<(f64, f64)> = pts0.iter().map(|&(x, y)| (x as f64, y as f64)).collect();
  let pts1f: Vec<(f64, f64)> = pts1.iter().map(|&(x, y)| (x as f64, y as f64)).collect();

  let mut rng = rand::rng();
  let mut best_inliers: Vec<usize> = Vec::new();
  let mut best_f: Option<Matrix3<f64>> = None;
  let mut iterations_needed = max_iterations;
  let mut iter = 0;

  while iter < iterations_needed {
    iter += 1;

    let sample = rand::seq::index::sample(&mut rng, n, 8);
    let sample0: Vec<(f64, f64)> = sample.iter().map(|i| pts0f[i]).collect();
    let sample1: Vec<(f64, f64)> = sample.iter().map(|i| pts1f[i]).collect();

    let Some(f) = fit_fundamental(&sample0, &sample1) else {
      continue;
    };

    let inliers: Vec<usize> = (0..n)
      .filter(|&i| sampson_distance(&f, pts0f[i], pts1f[i]) < inlier_threshold)
      .collect();

    if inliers.len() > best_inliers.len() {
      let inlier_ratio = inliers.len() as f64 / n as f64;
      best_inliers = inliers;
      best_f = Some(f);

      if inlier_ratio >= 1.0 {
        iterations_needed = iter;
      } else if inlier_ratio > 0.0 {
        let denom = (1.0 - inlier_ratio.powi(8)).ln();
        if denom < 0.0 {
          let needed = ((1.0 - confidence).ln() / denom).ceil();
          if needed.is_finite() {
            iterations_needed = iterations_needed.min((needed as usize).max(1));
          }
        }
      }
    }
  }

  let _ = best_f?;
  if best_inliers.len() < 8 {
    return None;
  }

  // Refine F using all inliers found so far, then recompute the inlier set against it.
  let in0: Vec<(f64, f64)> = best_inliers.iter().map(|&i| pts0f[i]).collect();
  let in1: Vec<(f64, f64)> = best_inliers.iter().map(|&i| pts1f[i]).collect();
  let refined_f = fit_fundamental(&in0, &in1)?;

  let final_inliers: Vec<usize> = (0..n)
    .filter(|&i| sampson_distance(&refined_f, pts0f[i], pts1f[i]) < inlier_threshold)
    .collect();

  let f_matrix: [f64; 9] = [
    refined_f[(0, 0)],
    refined_f[(0, 1)],
    refined_f[(0, 2)],
    refined_f[(1, 0)],
    refined_f[(1, 1)],
    refined_f[(1, 2)],
    refined_f[(2, 0)],
    refined_f[(2, 1)],
    refined_f[(2, 2)],
  ];

  Some(RansacResult {
    f_matrix,
    inliers: final_inliers,
  })
}

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

  /// Deterministic pseudo-random f64 in [0, 1), seeded by index -- avoids pulling in a
  /// distribution API just for a self-contained test.
  fn hash01(i: u32) -> f64 {
    let x = (i as f64) * 12.9898;
    (x.sin() * 43758.5453).fract().abs()
  }

  #[test]
  fn recovers_fundamental_matrix_from_synthetic_correspondences() {
    // Two pinhole cameras sharing intrinsics and orientation, camera1 translated
    // `baseline` along +x relative to camera0 (a simple rectified stereo rig). Points
    // projected from the same 3D point into both cameras are true epipolar
    // correspondences by construction, so a correct implementation should recover
    // nearly all of them as inliers.
    let fx = 800.0_f64;
    let fy = 800.0_f64;
    let cx = 320.0_f64;
    let cy = 240.0_f64;
    let baseline = 1.0_f64;

    let project = |(x, y, z): (f64, f64, f64), tx: f64| -> (f64, f64) {
      let x = x - tx;
      (fx * x / z + cx, fy * y / z + cy)
    };

    let points_3d: Vec<(f64, f64, f64)> = (0..200u32)
      .map(|i| {
        (
          hash01(i * 3) * 4.0 - 2.0,
          hash01(i * 3 + 1) * 4.0 - 2.0,
          5.0 + hash01(i * 3 + 2) * 5.0,
        )
      })
      .collect();

    let pts0: Vec<(f32, f32)> = points_3d
      .iter()
      .map(|&p| {
        let (x, y) = project(p, 0.0);
        (x as f32, y as f32)
      })
      .collect();
    let pts1: Vec<(f32, f32)> = points_3d
      .iter()
      .map(|&p| {
        let (x, y) = project(p, baseline);
        (x as f32, y as f32)
      })
      .collect();

    let result = ransac_fundamental_matrix(
      &pts0,
      &pts1,
      DEFAULT_MAX_ITERATIONS,
      DEFAULT_INLIER_THRESHOLD,
      DEFAULT_CONFIDENCE,
    )
    .expect("RANSAC should find a model for clean synthetic correspondences");

    assert!(
      result.inliers.len() as f64 >= 0.95 * pts0.len() as f64,
      "expected almost all {} synthetic correspondences to be inliers, got {}",
      pts0.len(),
      result.inliers.len()
    );
  }

  #[test]
  fn rejects_pairs_with_no_valid_geometry() {
    // Two independently-random point sets have no consistent epipolar relationship,
    // so RANSAC should never find enough inliers to pass MIN_CORRESPONDENCES.
    let pts0: Vec<(f32, f32)> = (0..200u32)
      .map(|i| ((hash01(i * 2) * 640.0) as f32, (hash01(i * 2 + 1) * 480.0) as f32))
      .collect();
    let pts1: Vec<(f32, f32)> = (0..200u32)
      .map(|i| {
        (
          (hash01(i * 2 + 1000) * 640.0) as f32,
          (hash01(i * 2 + 1001) * 480.0) as f32,
        )
      })
      .collect();

    let result = ransac_fundamental_matrix(
      &pts0,
      &pts1,
      DEFAULT_MAX_ITERATIONS,
      DEFAULT_INLIER_THRESHOLD,
      DEFAULT_CONFIDENCE,
    );
    if let Some(result) = result {
      assert!(
        result.inliers.len() < 20,
        "random, unrelated point sets should not produce a large inlier set, got {}",
        result.inliers.len()
      );
    }
  }
}