//! Pairwise geometry: a homography between two frames, found robustly.
//!
//! Two frames of a panorama are related by a rotation, and a rotation seen
//! through one lens is a homography of the image plane — `H = K R Kᵀ⁻¹`. The
//! homography is estimated first, from matches, because it does not need
//! the focal length; the focal length is then *read off* it (§ below), and
//! the rotation follows from both. This is the order Brown & Lowe (2007)
//! and OpenCV's stitcher use, and it is what makes the pipeline work when
//! EXIF says nothing about the lens.
//!
//! Coordinates throughout are **centred**: the principal point is the
//! origin. The focal formulae assume it, and centring before the DLT also
//! conditions the linear system — Hartley's normalisation, done once by the
//! caller rather than inside every solve.

use crate::linalg::{DMat, Mat3, Vec3};

/// A point in one image, centred on the principal point.
pub type Point = (f64, f64);

/// Apply a homography to a point.
pub fn apply(h: &Mat3, p: Point) -> Option<Point> {
    let v = *h * Vec3::new(p.0, p.1, 1.0);
    if v.z().abs() < 1e-12 {
        return None;
    }
    Some((v.x() / v.z(), v.y() / v.z()))
}

/// Least-squares homography from at least four correspondences by the
/// direct linear transform, with `h33` fixed at 1.
///
/// Fixing `h33` turns the homogeneous 8×9 system into an ordinary 8-unknown
/// least-squares problem that the normal equations and a Cholesky
/// factorisation solve without an SVD. The one homography it cannot
/// represent — `h33 = 0`, a point at the origin mapped to infinity — does
/// not occur between overlapping frames of one scene.
///
/// The points should be scaled to order one (divide by the focal length or
/// the image size) before calling: the normal equations square the
/// conditioning, and pixel coordinates in the thousands make them singular
/// in `f64`.
pub fn dlt(pairs: &[(Point, Point)]) -> Option<Mat3> {
    if pairs.len() < 4 {
        return None;
    }
    // Each pair gives two rows of A h = b with h = (h11..h32).
    //   x' = (h11 x + h12 y + h13) / (h31 x + h32 y + 1)
    //   → h11 x + h12 y + h13 - h31 x x' - h32 y x' = x'
    let mut ata = DMat::zeros(8);
    let mut atb = [0.0f64; 8];
    for &((x, y), (xp, yp)) in pairs {
        let rows: [([f64; 8], f64); 2] = [
            ([x, y, 1.0, 0.0, 0.0, 0.0, -x * xp, -y * xp], xp),
            ([0.0, 0.0, 0.0, x, y, 1.0, -x * yp, -y * yp], yp),
        ];
        for (a, b) in rows {
            for i in 0..8 {
                atb[i] += a[i] * b;
                for j in 0..8 {
                    ata[(i, j)] += a[i] * a[j];
                }
            }
        }
    }
    let h = ata.solve_spd(&atb)?;
    Some(Mat3([
        [h[0], h[1], h[2]],
        [h[3], h[4], h[5]],
        [h[6], h[7], 1.0],
    ]))
}

/// A homography with the correspondences that agree with it.
#[derive(Debug, Clone, PartialEq)]
pub struct RobustHomography {
    pub h: Mat3,
    /// Indices into the input pairs.
    pub inliers: Vec<usize>,
}

/// RANSAC over [`dlt`] on four-point samples, then a final least-squares
/// fit over every inlier.
///
/// `threshold` is the reprojection distance, in the same units as the
/// points, within which a pair counts as agreeing. The iteration count
/// adapts to the inlier ratio found so far in the usual way, capped at
/// `max_iterations`. `seed` makes a run reproducible (NFR-MRG-2): the
/// sampling is a small linear congruential generator, not the system's.
pub fn ransac_homography(
    pairs: &[(Point, Point)],
    threshold: f64,
    max_iterations: usize,
    seed: u64,
) -> Option<RobustHomography> {
    if pairs.len() < 4 {
        return None;
    }
    let n = pairs.len();
    let thr2 = threshold * threshold;
    let mut rng = Lcg(seed);
    let mut best: Option<(Vec<usize>, Mat3)> = None;
    let mut iterations = max_iterations;
    let mut i = 0;
    while i < iterations {
        i += 1;
        let sample = rng.distinct4(n);
        let Some(h) = dlt(&sample.map(|k| pairs[k])) else {
            continue;
        };
        let inliers: Vec<usize> = (0..n)
            .filter(|&k| agrees(&h, pairs[k], thr2))
            .collect();
        if best.as_ref().is_none_or(|(b, _)| inliers.len() > b.len()) {
            // Adapt: enough iterations to have drawn one all-inlier sample
            // with probability 0.99, given the ratio seen so far.
            let w = inliers.len() as f64 / n as f64;
            let p_all = w.powi(4);
            if p_all > 0.0 && p_all < 1.0 {
                let needed = ((1.0 - 0.99f64).ln() / (1.0 - p_all).ln()).ceil() as usize;
                iterations = iterations.min(needed.max(i + 1));
            }
            best = Some((inliers, h));
        }
    }
    let (inliers, h) = best?;
    if inliers.len() < 4 {
        return None;
    }
    // Refit on every inlier, and keep the refit only if it did not lose
    // support — a least-squares fit over a set with a few borderline points
    // can be pulled off the consensus the sample found.
    let refit: Vec<(Point, Point)> = inliers.iter().map(|&k| pairs[k]).collect();
    let h = match dlt(&refit) {
        Some(r) => {
            let count = (0..n).filter(|&k| agrees(&r, pairs[k], thr2)).count();
            if count >= inliers.len() {
                r
            } else {
                h
            }
        }
        None => h,
    };
    let inliers: Vec<usize> = (0..n).filter(|&k| agrees(&h, pairs[k], thr2)).collect();
    Some(RobustHomography { h, inliers })
}

fn agrees(h: &Mat3, (p, q): (Point, Point), thr2: f64) -> bool {
    match apply(h, p) {
        Some((x, y)) => {
            let (dx, dy) = (x - q.0, y - q.1);
            dx * dx + dy * dy <= thr2
        }
        None => false,
    }
}

/// The focal length a homography implies, if it implies one.
///
/// For `H = K R K⁻¹` with `K = diag(f, f, 1)` and the principal point at the
/// origin, the orthonormality of `R` gives two independent estimates of `f²`
/// from the first two rows and two from the first two columns; each is
/// taken where it is positive and the better-conditioned of the pair is
/// chosen, as OpenCV's `focalsFromHomography` does. The geometric mean of
/// the row and column estimates is returned. `None` when the homography is
/// too close to a pure translation to say anything — every estimate is then
/// a ratio of small numbers.
pub fn focal_from_homography(h: &Mat3) -> Option<f64> {
    let m = h.0;
    let (h00, h01, h02) = (m[0][0], m[0][1], m[0][2]);
    let (h10, h11, h12) = (m[1][0], m[1][1], m[1][2]);
    let (h20, h21) = (m[2][0], m[2][1]);

    let pick = |mut v1: f64, mut v2: f64, d1: f64, d2: f64| -> Option<f64> {
        if v1 < v2 {
            std::mem::swap(&mut v1, &mut v2);
        }
        if v1 > 0.0 && v2 > 0.0 {
            Some((if d1.abs() > d2.abs() { v1 } else { v2 }).sqrt())
        } else if v1 > 0.0 {
            Some(v1.sqrt())
        } else {
            None
        }
    };

    // From the third row.
    let d1 = h20 * h21;
    let d2 = (h21 - h20) * (h21 + h20);
    let f1 = if d1.abs() > 1e-12 || d2.abs() > 1e-12 {
        let v1 = if d1.abs() > 1e-12 {
            -(h00 * h01 + h10 * h11) / d1
        } else {
            f64::NAN
        };
        let v2 = if d2.abs() > 1e-12 {
            (h00 * h00 + h10 * h10 - h01 * h01 - h11 * h11) / d2
        } else {
            f64::NAN
        };
        pick(nan_to_neg(v1), nan_to_neg(v2), d1, d2)
    } else {
        None
    };

    // From the third column.
    let d1 = h00 * h10 + h01 * h11;
    let d2 = h00 * h00 + h01 * h01 - h10 * h10 - h11 * h11;
    let f0 = if d1.abs() > 1e-12 || d2.abs() > 1e-12 {
        let v1 = if d1.abs() > 1e-12 {
            -h02 * h12 / d1
        } else {
            f64::NAN
        };
        let v2 = if d2.abs() > 1e-12 {
            (h12 * h12 - h02 * h02) / d2
        } else {
            f64::NAN
        };
        pick(nan_to_neg(v1), nan_to_neg(v2), d1, d2)
    } else {
        None
    };

    match (f0, f1) {
        (Some(a), Some(b)) => Some((a * b).sqrt()),
        (Some(a), None) | (None, Some(a)) => Some(a),
        (None, None) => None,
    }
}

fn nan_to_neg(v: f64) -> f64 {
    if v.is_finite() {
        v
    } else {
        -1.0
    }
}

/// The rotation a homography encodes for a known focal length:
/// `R = K⁻¹ H K`, re-orthonormalised, with the scale of `H` divided out.
pub fn rotation_from_homography(h: &Mat3, f: f64) -> Mat3 {
    let m = h.0;
    // K⁻¹ H K with K = diag(f, f, 1): scale the third row by f and the
    // third column by 1/f.
    let r = Mat3([
        [m[0][0], m[0][1], m[0][2] / f],
        [m[1][0], m[1][1], m[1][2] / f],
        [m[2][0] * f, m[2][1] * f, m[2][2]],
    ]);
    r.orthonormalised()
}

/// A small deterministic generator for RANSAC's samples.
struct Lcg(u64);

impl Lcg {
    fn next(&mut self) -> u64 {
        // Knuth's MMIX constants.
        self.0 = self
            .0
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.0 >> 33
    }

    fn below(&mut self, n: usize) -> usize {
        (self.next() % n as u64) as usize
    }

    fn distinct4(&mut self, n: usize) -> [usize; 4] {
        let mut s = [0usize; 4];
        for i in 0..4 {
            loop {
                let k = self.below(n);
                if !s[..i].contains(&k) {
                    s[i] = k;
                    break;
                }
            }
        }
        s
    }
}

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

    /// Points under a known rotation seen through a known focal length,
    /// in centred image coordinates scaled by that focal length.
    fn synthetic(f: f64, r: Mat3, n: usize, noise: f64, seed: u64) -> Vec<(Point, Point)> {
        let mut rng = Lcg(seed);
        let mut out = Vec::new();
        while out.len() < n {
            // A point on the first image plane, within ±0.3 f of centre.
            let x = (rng.below(6001) as f64 - 3000.0) / 10000.0;
            let y = (rng.below(4001) as f64 - 2000.0) / 10000.0;
            let b = Vec3::new(x, y, 1.0);
            let v = r * b;
            if v.z() <= 0.2 {
                continue;
            }
            let nx = (rng.below(2001) as f64 - 1000.0) / 1000.0 * noise;
            let ny = (rng.below(2001) as f64 - 1000.0) / 1000.0 * noise;
            out.push(((x, y), (v.x() / v.z() + nx, v.y() / v.z() + ny)));
        }
        let _ = f;
        out
    }

    #[test]
    fn dlt_recovers_a_known_homography_exactly() {
        let r = Mat3::exp(Vec3::new(0.05, 0.3, 0.02));
        let pairs = synthetic(1.0, r, 12, 0.0, 1);
        let h = dlt(&pairs).expect("solvable");
        for &(p, q) in &pairs {
            let (x, y) = apply(&h, p).unwrap();
            assert!((x - q.0).abs() < 1e-9 && (y - q.1).abs() < 1e-9);
        }
    }

    #[test]
    fn ransac_finds_the_consensus_among_outliers() {
        let r = Mat3::exp(Vec3::new(-0.02, 0.25, 0.01));
        let mut pairs = synthetic(1.0, r, 60, 0.0005, 2);
        // Forty outliers: wrong second point.
        let mut rng = Lcg(9);
        for _ in 0..40 {
            let k = rng.below(60);
            let (p, _) = pairs[k];
            pairs.push((p, ((rng.below(1000) as f64 - 500.0) / 1000.0, 0.1)));
        }
        let robust = ransac_homography(&pairs, 0.003, 500, 3).expect("found");
        assert!(robust.inliers.len() >= 55, "{} inliers", robust.inliers.len());
        assert!(robust.inliers.iter().all(|&k| k < 60));
    }

    #[test]
    fn focal_is_read_off_a_rotation_homography() {
        // H in *pixel* coordinates for f = 1400: K R K⁻¹.
        let f = 1400.0;
        let r = Mat3::exp(Vec3::new(0.03, 0.35, -0.01));
        let m = r.0;
        let h = Mat3([
            [m[0][0], m[0][1], m[0][2] * f],
            [m[1][0], m[1][1], m[1][2] * f],
            [m[2][0] / f, m[2][1] / f, m[2][2]],
        ]);
        let est = focal_from_homography(&h).expect("estimable");
        assert!((est - f).abs() / f < 1e-6, "{est}");
        let back = rotation_from_homography(&h, f);
        for i in 0..3 {
            for j in 0..3 {
                assert!((back.0[i][j] - m[i][j]).abs() < 1e-9);
            }
        }
    }

    #[test]
    fn the_identity_implies_no_focal() {
        assert!(focal_from_homography(&Mat3::IDENTITY).is_none());
    }
}