//! TRACES: FR-MRG-1 | FR-MRG-5
//! From features to cameras: the alignment of a whole set.
//!
//! 1. Match every pair of frames (`matching`).
//! 2. For each pair with enough matches, a robust homography
//!    (`homography::ransac_homography`); a pair is a *link* when its inliers
//!    pass Brown & Lowe's test, `n_inliers > 8 + 0.3 · n_matches`, which
//!    is what separates a real overlap from a coincidence of descriptors.
//! 3. The focal length: the median of what the links' homographies imply,
//!    or the caller's hint if none of them implies anything.
//! 4. A spanning tree over the links, strongest first, from the
//!    best-connected frame; rotations chained along it.
//! 5. Bundle adjustment over every link's inliers (`bundle`).
//!
//! What it refuses to do is guess. A frame the tree does not reach is
//! reported by index with the reason (FR-MRG-5) and left out of the
//! cameras; the caller decides whether a set with a hole is worth
//! stitching, and the requirement says it is not.

use crate::bundle::{self, AdjustOptions, Cameras, Observation};
use crate::features::Features;
use crate::homography::{self, RobustHomography};
use crate::linalg::Mat3;
use crate::matching::{match_features, Match};
use crate::PanoError;

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AlignOptions {
    /// Descriptor similarity floor for a match (`matching`).
    pub min_similarity: f32,
    /// RANSAC agreement distance, in pixels of the features' image.
    pub ransac_px: f64,
    pub ransac_iterations: usize,
    /// A pair needs at least this many inliers to be a link, on top of
    /// Brown & Lowe's ratio test.
    pub min_inliers: usize,
    /// Focal length in pixels of the features' image, if the caller knows
    /// it (EXIF and a sensor width). Used only when the homographies do not
    /// determine one.
    pub focal_hint: Option<f64>,
    pub adjust: AdjustOptions,
    /// For RANSAC's sampling: the same seed gives the same alignment
    /// (NFR-MRG-2).
    pub seed: u64,
}

impl Default for AlignOptions {
    fn default() -> Self {
        AlignOptions {
            min_similarity: 0.82,
            ransac_px: 3.0,
            ransac_iterations: 1000,
            min_inliers: 12,
            focal_hint: None,
            adjust: AdjustOptions::default(),
            seed: 0x5eed,
        }
    }
}

/// An overlap the alignment trusts.
#[derive(Debug, Clone, PartialEq)]
pub struct Link {
    pub i: usize,
    pub j: usize,
    pub matches: usize,
    pub inliers: usize,
    /// Maps centred points of `i` to centred points of `j`.
    pub h: Mat3,
}

/// Why a frame is not in the alignment.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Unaligned {
    /// Not enough matches with any other frame to try a geometry.
    NoMatches,
    /// Matches existed but none survived RANSAC as a real overlap.
    NoOverlap,
    /// Overlaps existed but only with frames that are themselves unaligned.
    Disconnected,
}

impl std::fmt::Display for Unaligned {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Unaligned::NoMatches => "too few matching features with any other frame",
            Unaligned::NoOverlap => "no consistent overlap with any other frame",
            Unaligned::Disconnected => "overlaps only with frames that could not be aligned",
        })
    }
}

/// The result: cameras for the aligned frames, and the rest named.
#[derive(Debug, Clone, PartialEq)]
pub struct Alignment {
    /// One rotation per input frame, camera to world, for aligned frames;
    /// `None` for the unaligned. The reference frame is the best-connected
    /// one and has the identity.
    pub rotations: Vec<Option<Mat3>>,
    /// Focal length in pixels of the features' image.
    pub focal: f64,
    pub links: Vec<Link>,
    pub unaligned: Vec<(usize, Unaligned)>,
    /// Bundle adjustment's RMS reprojection error, in pixels.
    pub rms_px: f64,
}

impl Alignment {
    pub fn is_complete(&self) -> bool {
        self.unaligned.is_empty()
    }

    /// The cameras of the aligned frames, indexed as the input — a frame
    /// that is not aligned is given the identity, so this is only useful
    /// when [`Self::is_complete`].
    pub fn cameras(&self) -> Cameras {
        Cameras {
            rotations: self
                .rotations
                .iter()
                .map(|r| r.unwrap_or(Mat3::IDENTITY))
                .collect(),
            focal: self.focal,
        }
    }
}

/// Align a set of frames from their features.
///
/// Every `Features` must be in its own frame's pixel coordinates with the
/// image size filled in; points are centred on the image centre here. The
/// frames must all come from the same lens at the same focal length, which
/// is the panorama assumption and not checked — the caller has the EXIF.
pub fn align(frames: &[Features], opts: &AlignOptions) -> Result<Alignment, PanoError> {
    let n = frames.len();
    if n < 2 {
        return Err(PanoError::Input(
            "a panorama needs at least two frames".into(),
        ));
    }

    let centre = |k: usize, i: usize| -> (f64, f64) {
        let kp = frames[k].keypoints[i];
        (
            f64::from(kp.x) - frames[k].width as f64 / 2.0,
            f64::from(kp.y) - frames[k].height as f64 / 2.0,
        )
    };
    // Scale for the DLT's conditioning: points of order one.
    let scale = 1.0
        / frames
            .iter()
            .map(|f| f.width.max(f.height) as f64)
            .fold(1.0, f64::max);

    // 1 + 2: every pair.
    let mut links = Vec::new();
    let mut observations: Vec<Observation> = Vec::new();
    let mut matched_any = vec![false; n];
    let t_match = std::time::Instant::now();
    for i in 0..n {
        for j in i + 1..n {
            let matches: Vec<Match> = match_features(&frames[i], &frames[j], opts.min_similarity);
            log::debug!("pair {i}-{j}: {} matches", matches.len());
            if matches.len() < 4 {
                continue;
            }
            matched_any[i] = true;
            matched_any[j] = true;
            let pairs: Vec<((f64, f64), (f64, f64))> = matches
                .iter()
                .map(|m| {
                    let (a, b) = (centre(i, m.a), centre(j, m.b));
                    ((a.0 * scale, a.1 * scale), (b.0 * scale, b.1 * scale))
                })
                .collect();
            let Some(RobustHomography { h, inliers }) = homography::ransac_homography(
                &pairs,
                opts.ransac_px * scale,
                opts.ransac_iterations,
                opts.seed ^ ((i as u64) << 32 | j as u64),
            ) else {
                continue;
            };
            let needed = (8.0 + 0.3 * matches.len() as f64).ceil() as usize;
            log::debug!("pair {i}-{j}: {} inliers, {needed} needed", inliers.len());
            if inliers.len() <= needed || inliers.len() < opts.min_inliers {
                continue;
            }
            // Back to pixels: H_px = S⁻¹ H S.
            let m = h.0;
            let h_px = Mat3([
                [m[0][0], m[0][1], m[0][2] / scale],
                [m[1][0], m[1][1], m[1][2] / scale],
                [m[2][0] * scale, m[2][1] * scale, m[2][2]],
            ]);
            for &k in &inliers {
                let (a, b) = pairs[k];
                observations.push(Observation {
                    i,
                    j,
                    pi: (a.0 / scale, a.1 / scale),
                    pj: (b.0 / scale, b.1 / scale),
                });
            }
            links.push(Link {
                i,
                j,
                matches: matches.len(),
                inliers: inliers.len(),
                h: h_px,
            });
        }
    }

    log::debug!("matching and pairwise geometry in {:?}", t_match.elapsed());

    // 3: the focal length.
    let mut estimates: Vec<f64> = links
        .iter()
        .filter_map(|l| homography::focal_from_homography(&l.h))
        .filter(|f| f.is_finite() && *f > 0.0)
        .collect();
    let longest = frames
        .iter()
        .map(|f| f.width.max(f.height) as f64)
        .fold(0.0, f64::max);
    let focal = if !estimates.is_empty() {
        estimates.sort_by(f64::total_cmp);
        let median = estimates[estimates.len() / 2];
        // A homography of a nearly pure pan can imply almost anything;
        // clamp to the range a real lens on this sensor can reach.
        median.clamp(0.3 * longest, 6.0 * longest)
    } else if let Some(hint) = opts.focal_hint {
        hint
    } else {
        // No overlap said anything and nobody told us: a normal lens.
        longest
    };

    // 4: spanning tree, strongest link first, from the best-connected frame.
    let mut rotations: Vec<Option<Mat3>> = vec![None; n];
    let mut unaligned = Vec::new();
    if links.is_empty() {
        for (k, &matched) in matched_any.iter().enumerate() {
            unaligned.push((
                k,
                if matched {
                    Unaligned::NoOverlap
                } else {
                    Unaligned::NoMatches
                },
            ));
        }
        return Ok(Alignment {
            rotations,
            focal,
            links,
            unaligned,
            rms_px: 0.0,
        });
    }
    let mut degree = vec![0usize; n];
    for l in &links {
        degree[l.i] += l.inliers;
        degree[l.j] += l.inliers;
    }
    let root = (0..n).max_by_key(|&k| degree[k]).unwrap_or(0);
    rotations[root] = Some(Mat3::IDENTITY);
    loop {
        // The strongest link from an aligned frame to an unaligned one.
        let best = links
            .iter()
            .filter(|l| rotations[l.i].is_some() != rotations[l.j].is_some())
            .max_by_key(|l| l.inliers);
        let Some(l) = best else { break };
        let r_ij = homography::rotation_from_homography(&l.h, focal);
        // H_ij takes points of i to j, so bearings b_j = R_ij b_i, and with
        // world = R_i · cam_i: R_j = R_i · R_ijᵀ.
        if let Some(ri) = rotations[l.i] {
            rotations[l.j] = Some((ri * r_ij.transpose()).orthonormalised());
        } else if let Some(rj) = rotations[l.j] {
            rotations[l.i] = Some((rj * r_ij).orthonormalised());
        }
    }
    for k in 0..n {
        if rotations[k].is_none() {
            let reason = if !matched_any[k] {
                Unaligned::NoMatches
            } else if links.iter().any(|l| l.i == k || l.j == k) {
                Unaligned::Disconnected
            } else {
                Unaligned::NoOverlap
            };
            unaligned.push((k, reason));
        }
    }

    // 5: adjust the aligned frames together. The reference frame must be
    // index 0 of the adjustment (it holds frame 0 fixed), so the aligned
    // frames are renumbered with the root first.
    let aligned: Vec<usize> = std::iter::once(root)
        .chain((0..n).filter(|&k| k != root && rotations[k].is_some()))
        .collect();
    let index_of = |k: usize| aligned.iter().position(|&a| a == k);
    let start = Cameras {
        rotations: aligned.iter().map(|&k| rotations[k].unwrap()).collect(),
        focal,
    };
    let obs: Vec<Observation> = observations
        .iter()
        .filter_map(|o| {
            Some(Observation {
                i: index_of(o.i)?,
                j: index_of(o.j)?,
                pi: o.pi,
                pj: o.pj,
            })
        })
        .collect();
    let t_adjust = std::time::Instant::now();
    let adjusted = bundle::adjust(start, &obs, &opts.adjust)?;
    log::debug!(
        "bundle adjustment: {} observations, {} iterations in {:?}",
        obs.len(),
        adjusted.iterations,
        t_adjust.elapsed()
    );
    for (slot, &k) in aligned.iter().enumerate() {
        rotations[k] = Some(adjusted.cameras.rotations[slot]);
    }

    Ok(Alignment {
        rotations,
        focal: adjusted.cameras.focal,
        links,
        unaligned,
        rms_px: adjusted.rms_px,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::features::{Keypoint, DESCRIPTOR_LEN};
    use crate::linalg::Vec3;

    /// Frames of a synthetic sweep: world directions with random unit
    /// descriptors, each frame seeing the ones in its field of view.
    fn synthetic_sweep(
        n: usize,
        step: f64,
        f: f64,
        w: usize,
        h: usize,
    ) -> (Vec<Features>, Cameras) {
        let mut seed = 777u64;
        let mut rnd = || {
            seed = seed
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((seed >> 33) as f64 / (1u64 << 31) as f64) - 0.5
        };
        let rotations: Vec<Mat3> = (0..n)
            .map(|k| {
                Mat3::rotation(Vec3::new(0.0, 1.0, 0.0), step * k as f64)
                    * Mat3::rotation(Vec3::new(1.0, 0.0, 0.0), 0.02 * ((k % 3) as f64 - 1.0))
            })
            .collect();
        let truth = Cameras {
            rotations,
            focal: f,
        };
        let total = step * (n as f64 - 1.0);
        let mut frames: Vec<Features> = (0..n)
            .map(|_| Features {
                keypoints: Vec::new(),
                descriptors: Vec::new(),
                width: w,
                height: h,
            })
            .collect();
        for _ in 0..600 * 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(),
            );
            let desc: Vec<f32> = (0..DESCRIPTOR_LEN).map(|_| rnd() as f32).collect();
            let norm = desc.iter().map(|v| v * v).sum::<f32>().sqrt();
            let desc: Vec<f32> = desc.iter().map(|v| v / norm).collect();
            for (k, frame) in frames.iter_mut().enumerate() {
                if let Some(p) = truth.project(k, d) {
                    let (x, y) = (p.0 + w as f64 / 2.0, p.1 + h as f64 / 2.0);
                    if x >= 0.0 && x < w as f64 && y >= 0.0 && y < h as f64 {
                        frame.keypoints.push(Keypoint {
                            x: (x + rnd() * 0.6) as f32,
                            y: (y + rnd() * 0.6) as f32,
                            score: 1.0,
                        });
                        frame.descriptors.extend_from_slice(&desc);
                    }
                }
            }
        }
        (frames, truth)
    }

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

    #[test]
    fn a_synthetic_sweep_is_aligned_to_its_truth() {
        let (frames, truth) = synthetic_sweep(6, 0.3, 1400.0, 1024, 768);
        let out = align(&frames, &AlignOptions::default()).expect("aligned");
        assert!(out.is_complete(), "unaligned: {:?}", out.unaligned);
        assert_eq!(out.links.len(), 5 + 4, "links: {}", out.links.len());
        assert!((out.focal - 1400.0).abs() < 15.0, "focal {}", out.focal);
        assert!(out.rms_px < 1.0, "rms {}", out.rms_px);
        // Relative rotations match the truth's, whichever frame is the root.
        let root = out
            .rotations
            .iter()
            .position(|r| *r == Some(Mat3::IDENTITY))
            .unwrap();
        for k in 0..6 {
            let rel_truth = truth.rotations[root].transpose() * truth.rotations[k];
            let rel_out = out.rotations[k].unwrap();
            let err = angle_between(rel_truth, rel_out);
            assert!(err < 2e-3, "frame {k} off by {err} rad");
        }
    }

    #[test]
    fn a_frame_from_nowhere_is_named_not_guessed() {
        let (mut frames, _) = synthetic_sweep(4, 0.3, 1400.0, 1024, 768);
        // Frame 3 gets descriptors nobody else has.
        for v in &mut frames[3].descriptors {
            *v = -*v;
        }
        let out = align(&frames, &AlignOptions::default()).expect("aligned");
        assert_eq!(out.unaligned.len(), 1);
        assert_eq!(out.unaligned[0].0, 3);
        assert!(out.rotations[3].is_none());
        assert!(out.rotations[..3].iter().all(Option::is_some));
    }

    #[test]
    fn one_frame_is_refused() {
        let (frames, _) = synthetic_sweep(1, 0.3, 1400.0, 640, 480);
        assert!(matches!(
            align(&frames, &AlignOptions::default()),
            Err(PanoError::Input(_))
        ));
    }
}