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;

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 < 8 || 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 < 8 || 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,
  })
}