//! Bundle adjustment: every rotation and the focal length, refined together.
//!
//! The pairwise homographies (`homography.rs`) each know about two frames.
//! Chained around a loop they disagree with themselves by the accumulated
//! error, and a twelve-frame sweep chained end to end drifts by a visible
//! amount. This solves for all the rotations at once, against every inlier
//! match of every pair, so the error is spread rather than accumulated —
//! Brown & Lowe's step 4, with the camera model reduced to what a panorama
//! needs: one rotation per frame and one focal length shared by all.
//!
//! Levenberg–Marquardt with a numerical Jacobian. Analytic derivatives of a
//! rotation's projection are not hard, but they are a second place the
//! model is written down, and the model is small: forty parameters, a few
//! thousand residuals, a Jacobian that costs forty residual evaluations.
//! The whole solve is milliseconds. Correctness over cleverness, and one
//! definition of the projection to keep right.

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

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

/// One inlier match between two frames.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Observation {
    pub i: usize,
    pub j: usize,
    pub pi: Point,
    pub pj: Point,
}

/// What the adjustment starts from and returns: a rotation per frame
/// (camera to world; frame 0 is the world) and the focal length in pixels.
#[derive(Debug, Clone, PartialEq)]
pub struct Cameras {
    pub rotations: Vec<Mat3>,
    pub focal: f64,
}

impl Cameras {
    /// The unit direction, in world space, that pixel `p` of frame `i` looks
    /// along.
    pub fn bearing(&self, i: usize, p: Point) -> Vec3 {
        self.rotations[i] * Vec3::new(p.0, p.1, self.focal).normalised()
    }

    /// Where world direction `d` lands in frame `j`, or `None` if it is
    /// behind the camera.
    pub fn project(&self, j: usize, d: Vec3) -> Option<Point> {
        let c = self.rotations[j].transpose() * d;
        if c.z() <= 1e-9 {
            return None;
        }
        Some((self.focal * c.x() / c.z(), self.focal * c.y() / c.z()))
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AdjustOptions {
    pub max_iterations: usize,
    /// Residuals beyond this many pixels are down-weighted (Huber), so a
    /// mismatch RANSAC let through pulls with bounded force.
    pub huber_px: f64,
    /// Whether the focal length is a free parameter. Off, it is held at the
    /// starting value — for a set whose rotations are all small, the focal
    /// length is weakly observable and better taken from the homographies'
    /// median than pulled about by noise.
    pub refine_focal: bool,
}

impl Default for AdjustOptions {
    fn default() -> Self {
        AdjustOptions {
            max_iterations: 50,
            huber_px: 3.0,
            refine_focal: true,
        }
    }
}

/// The adjusted cameras and the fit.
#[derive(Debug, Clone, PartialEq)]
pub struct Adjusted {
    pub cameras: Cameras,
    /// Root-mean-square reprojection error over all observations, in pixels
    /// (unweighted, so an outlier RANSAC missed shows here rather than
    /// hiding under its Huber weight).
    pub rms_px: f64,
    pub iterations: usize,
}

/// Refine `start` against `observations`.
///
/// Frame 0's rotation is held fixed: the world frame is arbitrary and
/// fixing one camera removes the freedom. Every other frame must appear in
/// at least one observation or its rotation is undetermined and the normal
/// equations are singular — the caller (`align`) guarantees it by only
/// adjusting frames a spanning tree reached.
pub fn adjust(
    start: Cameras,
    observations: &[Observation],
    opts: &AdjustOptions,
) -> Result<Adjusted, PanoError> {
    let n_frames = start.rotations.len();
    if n_frames < 2 || observations.is_empty() {
        let rms = rms(&start, observations);
        return Ok(Adjusted {
            cameras: start,
            rms_px: rms,
            iterations: 0,
        });
    }
    // Every adjustable frame must be constrained by something, or its
    // block of the normal equations is zero and the solve is meaningless —
    // checked here, by name, rather than left to surface as a step that
    // fails to lower the cost.
    let mut seen = vec![false; n_frames];
    for o in observations {
        seen[o.i] = true;
        seen[o.j] = true;
    }
    if let Some(k) = (1..n_frames).find(|&k| !seen[k]) {
        return Err(PanoError::Geometry(format!(
            "frame {k} has no observations constraining it"
        )));
    }
    let n_rot = 3 * (n_frames - 1);
    let n_params = n_rot + usize::from(opts.refine_focal);
    let n_res = 2 * observations.len();

    // Parameters are *increments* on the current cameras, re-applied each
    // accepted step: rotation k ← exp(δ_k) · rotation k, focal ← f · exp(δ_f).
    // Composing on the left keeps the increment in world space, where a
    // small rotation means the same thing for every frame.
    let apply = |base: &Cameras, x: &[f64]| -> Cameras {
        let mut rotations = base.rotations.clone();
        for k in 1..n_frames {
            let w = Vec3::new(x[3 * (k - 1)], x[3 * (k - 1) + 1], x[3 * (k - 1) + 2]);
            rotations[k] = (Mat3::exp(w) * base.rotations[k]).orthonormalised();
        }
        let focal = if opts.refine_focal {
            base.focal * x[n_rot].exp()
        } else {
            base.focal
        };
        Cameras { rotations, focal }
    };

    let residuals = |c: &Cameras, out: &mut Vec<f64>| {
        out.clear();
        for o in observations {
            let d = c.bearing(o.i, o.pi);
            match c.project(o.j, d) {
                Some((x, y)) => {
                    out.push(x - o.pj.0);
                    out.push(y - o.pj.1);
                }
                None => {
                    // Behind the camera: as wrong as a residual can be
                    // without being infinite. The Huber weight caps its pull.
                    out.push(1e4);
                    out.push(1e4);
                }
            }
        }
    };

    let weights = |r: &[f64], out: &mut Vec<f64>| {
        out.clear();
        for pair in r.chunks_exact(2) {
            let m = (pair[0] * pair[0] + pair[1] * pair[1]).sqrt();
            let w = if m > opts.huber_px {
                opts.huber_px / m
            } else {
                1.0
            };
            out.push(w);
            out.push(w);
        }
    };

    // The robust cost itself, not the weighted sum of squares: the weights
    // above are the IRLS linearisation for one step, and comparing two
    // steps by sums taken under different weights would accept the wrong
    // ones. Huber: quadratic within the threshold, linear beyond it.
    let cost = |r: &[f64]| -> f64 {
        r.chunks_exact(2)
            .map(|pair| {
                let m = (pair[0] * pair[0] + pair[1] * pair[1]).sqrt();
                if m <= opts.huber_px {
                    m * m
                } else {
                    2.0 * opts.huber_px * m - opts.huber_px * opts.huber_px
                }
            })
            .sum()
    };

    let mut cameras = start;
    let mut r = Vec::with_capacity(n_res);
    let mut w = Vec::with_capacity(n_res);
    residuals(&cameras, &mut r);
    weights(&r, &mut w);
    let mut current = cost(&r);

    let mut lambda = 1e-3;
    let mut jac = vec![0.0f64; n_res * n_params];
    let mut r_plus = Vec::with_capacity(n_res);
    let zero = vec![0.0f64; n_params];
    let mut iterations = 0;

    for _ in 0..opts.max_iterations {
        iterations += 1;

        // Numerical Jacobian about the current cameras (x = 0).
        const H: f64 = 1e-6;
        for p in 0..n_params {
            let mut x = zero.clone();
            x[p] = H;
            let c_plus = apply(&cameras, &x);
            residuals(&c_plus, &mut r_plus);
            for (k, (rp, r0)) in r_plus.iter().zip(&r).enumerate() {
                jac[k * n_params + p] = (rp - r0) / H;
            }
        }

        // Normal equations, weighted: (JᵀWJ + λ·diag) δ = −JᵀWr.
        let mut a = DMat::zeros(n_params);
        let mut b = vec![0.0f64; n_params];
        for k in 0..n_res {
            let row = &jac[k * n_params..(k + 1) * n_params];
            let wk = w[k];
            for p in 0..n_params {
                b[p] -= wk * row[p] * r[k];
                for q in 0..n_params {
                    a[(p, q)] += wk * row[p] * row[q];
                }
            }
        }

        // Try steps with increasing damping until one lowers the cost.
        let mut accepted = false;
        for _ in 0..10 {
            let mut damped = a.clone();
            for p in 0..n_params {
                let d = a[(p, p)];
                damped[(p, p)] = d + lambda * d.max(1e-9);
            }
            let Some(delta) = damped.solve_spd(&b) else {
                return Err(PanoError::Geometry(
                    "the adjustment's normal equations are singular: a frame has no \
                     observations constraining it"
                        .into(),
                ));
            };
            let candidate = apply(&cameras, &delta);
            residuals(&candidate, &mut r_plus);
            let c_new = cost(&r_plus);
            if c_new < current {
                let improvement = (current - c_new) / current.max(1e-12);
                let step: f64 = delta.iter().map(|d| d * d).sum::<f64>().sqrt();
                cameras = candidate;
                std::mem::swap(&mut r, &mut r_plus);
                weights(&r, &mut w);
                current = c_new;
                lambda = (lambda / 3.0).max(1e-9);
                accepted = true;
                // Converged when a *lightly damped* step no longer helps. A
                // heavily damped step is small by construction and would
                // pass an improvement test long before the minimum.
                if step < 1e-10 || (improvement < 1e-8 && lambda < 1e-2) {
                    return Ok(Adjusted {
                        rms_px: rms(&cameras, observations),
                        cameras,
                        iterations,
                    });
                }
                break;
            }
            lambda *= 5.0;
        }
        if !accepted {
            break;
        }
    }

    Ok(Adjusted {
        rms_px: rms(&cameras, observations),
        cameras,
        iterations,
    })
}

/// Unweighted RMS reprojection error in pixels.
pub fn rms(c: &Cameras, observations: &[Observation]) -> f64 {
    if observations.is_empty() {
        return 0.0;
    }
    let sum: f64 = observations
        .iter()
        .map(|o| match c.project(o.j, c.bearing(o.i, o.pi)) {
            Some((x, y)) => (x - o.pj.0).powi(2) + (y - o.pj.1).powi(2),
            None => 1e8,
        })
        .sum();
    (sum / observations.len() as f64).sqrt()
}

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

    /// A synthetic sweep: `n` cameras panned by `step` radians each with a
    /// little pitch and roll, `f` pixels, and matches between neighbours
    /// from a cloud of world directions.
    fn sweep(n: usize, step: f64, f: f64, noise_px: f64) -> (Cameras, Vec<Observation>) {
        let mut rotations = Vec::new();
        for k in 0..n {
            let yaw = step * k as f64;
            let pitch = 0.01 * ((k * 7) % 3) as f64;
            let roll = 0.005 * ((k * 5) % 4) as f64;
            let r = Mat3::rotation(Vec3::new(0.0, 1.0, 0.0), yaw)
                * Mat3::rotation(Vec3::new(1.0, 0.0, 0.0), pitch)
                * Mat3::rotation(Vec3::new(0.0, 0.0, 1.0), roll);
            rotations.push(r);
        }
        let truth = Cameras {
            rotations,
            focal: f,
        };

        // World directions: a fan across the whole sweep.
        let mut obs = Vec::new();
        let mut seed = 12345u64;
        let mut rnd = || {
            seed = seed
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((seed >> 33) as f64 / (1u64 << 31) as f64) - 0.5
        };
        let total = step * (n as f64 - 1.0);
        for _ in 0..400 * n {
            let yaw = rnd() * (total + 0.8) + total / 2.0;
            let pitch = rnd() * 0.5;
            let d = Vec3::new(
                yaw.sin() * pitch.cos(),
                pitch.sin(),
                yaw.cos() * pitch.cos(),
            )
            .normalised();
            // Visible in which frames? Within ±0.35 f of centre.
            let mut seen: Vec<(usize, Point)> = Vec::new();
            for k in 0..n {
                if let Some(p) = truth.project(k, d) {
                    if p.0.abs() < 0.35 * f && p.1.abs() < 0.25 * f {
                        seen.push((k, (p.0 + rnd() * noise_px, p.1 + rnd() * noise_px)));
                    }
                }
            }
            for a in 0..seen.len() {
                for b in a + 1..seen.len() {
                    obs.push(Observation {
                        i: seen[a].0,
                        j: seen[b].0,
                        pi: seen[a].1,
                        pj: seen[b].1,
                    });
                }
            }
        }
        (truth, obs)
    }

    fn angle_between(a: Mat3, b: Mat3) -> f64 {
        (a.transpose() * b).log().norm()
    }

    #[test]
    fn a_perturbed_start_converges_back_to_the_truth() {
        let (truth, obs) = sweep(6, 0.3, 1400.0, 0.0);
        assert!(obs.len() > 500);
        // Perturb every rotation but the first by ~1°, and the focal by 5%.
        let mut start = truth.clone();
        for k in 1..6 {
            let w = Vec3::new(0.01, -0.015, 0.008) * (k as f64 / 3.0);
            start.rotations[k] = Mat3::exp(w) * start.rotations[k];
        }
        start.focal *= 1.05;
        let before = rms(&start, &obs);
        let out = adjust(start, &obs, &AdjustOptions::default()).expect("solvable");
        assert!(out.rms_px < 1e-3, "rms {} (was {before})", out.rms_px);
        assert!(
            (out.cameras.focal - 1400.0).abs() < 0.5,
            "focal {}",
            out.cameras.focal
        );
        for k in 0..6 {
            let err = angle_between(out.cameras.rotations[k], truth.rotations[k]);
            assert!(err < 1e-5, "frame {k} off by {err} rad");
        }
    }

    #[test]
    fn noise_is_averaged_rather_than_accumulated() {
        let (truth, obs) = sweep(8, 0.25, 1400.0, 1.0);
        let mut start = truth.clone();
        for k in 1..8 {
            start.rotations[k] =
                Mat3::exp(Vec3::new(0.0, 0.004 * k as f64, 0.0)) * start.rotations[k];
        }
        let out = adjust(start, &obs, &AdjustOptions::default()).expect("solvable");
        // ±0.5 px of uniform noise on every coordinate has an RMS of 0.41 px
        // per axis, so the fit's RMS over both axes should sit near 0.58 and
        // cannot be much below it.
        assert!(out.rms_px < 0.7, "rms {}", out.rms_px);

        // The focal length and the sweep are nearly degenerate for a
        // single row: only the perspective inside each overlap pins the
        // focal, and a pixel of noise is worth about a tenth of a percent of
        // it. What that error does is scale every yaw by the same factor —
        // a uniform stretch of the panorama, invisible in the result — so the
        // absolute rotation error grows linearly along the sweep and is not
        // the measure of the solve. The residual after removing that stretch
        // is.
        let f_ratio = out.cameras.focal / 1400.0;
        assert!((f_ratio - 1.0).abs() < 5e-3, "focal {}", out.cameras.focal);
        for k in 0..8 {
            let yaw_k = 0.25 * k as f64;
            let expected_stretch = (f_ratio - 1.0).abs() * yaw_k;
            let err = angle_between(out.cameras.rotations[k], truth.rotations[k]);
            assert!(
                err < expected_stretch + 1.5e-4,
                "frame {k} off by {err} rad, {expected_stretch} of it the focal's"
            );
        }
    }

    #[test]
    fn a_frame_without_observations_is_refused() {
        let (truth, mut obs) = sweep(4, 0.3, 1400.0, 0.0);
        obs.retain(|o| o.i != 3 && o.j != 3);
        let err = adjust(truth, &obs, &AdjustOptions::default()).unwrap_err();
        assert!(matches!(err, PanoError::Geometry(_)));
    }
}