projective-grid 0.10.1

//! Shared output types for detection and consistency tasks.
//!
//! The detection surface is pinned to `f32`; see [`crate::feature`] for
//! the rationale.

use std::collections::HashMap;

use nalgebra::{Point2, Projective2};

use crate::lattice::{Coord, GridDimensions, LatticeKind};

/// One labelled grid feature in a solved grid.
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub struct GridEntry {
    /// Lattice coordinate assigned to this feature.
    pub coord: Coord,
    /// Caller-owned feature source index.
    pub source_index: usize,
    /// Image-frame pixel-center position.
    pub image_position: Point2<f32>,
    /// Reprojection residual in image pixels, when a fit was computed.
    pub residual_px: Option<f32>,
}

impl GridEntry {
    /// Construct a labelled grid entry.
    pub fn new(
        coord: Coord,
        source_index: usize,
        image_position: Point2<f32>,
        residual_px: Option<f32>,
    ) -> Self {
        Self {
            coord,
            source_index,
            image_position,
            residual_px,
        }
    }
}

/// A labelled grid component.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct LabelledGrid {
    /// Lattice family of this grid.
    pub lattice: LatticeKind,
    /// Labelled feature entries.
    pub entries: Vec<GridEntry>,
    /// Inclusive coordinate bounding box, if the grid is non-empty.
    pub bbox: Option<(Coord, Coord)>,
    /// Optional known dimensions supplied by the caller.
    pub dimensions: Option<GridDimensions>,
}

impl LabelledGrid {
    /// Construct a labelled grid.
    pub fn new(
        lattice: LatticeKind,
        entries: Vec<GridEntry>,
        dimensions: Option<GridDimensions>,
    ) -> Self {
        let bbox = bbox_for_entries(&entries);
        Self {
            lattice,
            entries,
            bbox,
            dimensions,
        }
    }

    /// Linear-scan lookup of the labelled entry with the given source index.
    pub fn find(&self, source_index: usize) -> Option<&GridEntry> {
        self.entries.iter().find(|e| e.source_index == source_index)
    }

    /// Normalize the labelled grid in place to the canonical output frame.
    ///
    /// Three steps, in order:
    ///
    /// 1. **Rebase** the coordinate bounding-box minimum to `(0, 0)`, so every
    ///    `coord` is non-negative (the hard non-negative-label invariant for
    ///    overlay / calibration consumers).
    /// 2. **Canonicalize orientation** so the first lattice axis (`u`) points
    ///    roughly `+x` (right) and the second (`v`) roughly `+y` (down) in image
    ///    pixels. The grid finder assigns `(u, v)` from its internal axis-slot
    ///    convention, which has no relation to image orientation; without this
    ///    step `(0, 0)` can land anywhere on the detected grid. The decision is
    ///    driven by [`GridEntry::image_position`] (averaged step vectors over all
    ///    adjacent labelled pairs); positions are never modified, only labels are
    ///    permuted / sign-flipped. When this step transposes the two axes, any
    ///    caller-supplied [`dimensions`](LabelledGrid::dimensions) have their
    ///    `width`/`height` swapped too, so they stay aligned with the new axes.
    /// 3. **Sort** entries by `(v, u)` for a stable output order, and recompute
    ///    [`bbox`](LabelledGrid::bbox).
    ///
    /// This is the single source of truth for grid-result normalization: target
    /// detectors call it instead of re-implementing rebase / canonicalize / sort
    /// at their output stage. It operates only on the labelled grid, so any
    /// [`LatticeFit`] computed against the *pre*-normalization labels is no
    /// longer valid afterwards — normalize before fitting, or refit.
    pub fn normalize(&mut self) {
        rebase_entries_to_origin(&mut self.entries);
        let swapped = canonicalize_to_image_axes(&mut self.entries);
        if swapped {
            // The `u` ↔ `v` transpose must carry to caller-supplied dimensions,
            // or `dimensions` would describe the pre-swap axes while `entries`
            // and `bbox` describe the new ones.
            if let Some(dims) = self.dimensions.as_mut() {
                std::mem::swap(&mut dims.width, &mut dims.height);
            }
        }
        self.entries.sort_by_key(|e| (e.coord.v, e.coord.u));
        self.bbox = bbox_for_entries(&self.entries);
    }
}

/// Residual summary in image pixels.
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub struct ResidualSummary {
    /// Number of residuals included in the summary.
    pub count: usize,
    /// Mean residual in pixels.
    pub mean_px: f32,
    /// Maximum residual in pixels.
    pub max_px: f32,
}

impl ResidualSummary {
    /// Construct a residual summary.
    pub fn new(count: usize, mean_px: f32, max_px: f32) -> Self {
        Self {
            count,
            mean_px,
            max_px,
        }
    }
}

/// Fitted lattice-to-image transform plus residual summary.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct LatticeFit {
    /// Projective mapping from model-plane lattice coordinates to image pixels.
    pub model_to_image: Projective2<f32>,
    /// Residual summary in image pixels.
    pub residuals: ResidualSummary,
}

impl LatticeFit {
    /// Construct a lattice fit.
    pub fn new(model_to_image: Projective2<f32>, residuals: ResidualSummary) -> Self {
        Self {
            model_to_image,
            residuals,
        }
    }
}

/// Reason why an observed feature did not pass a task gate.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum RejectionReason {
    /// Reprojection residual exceeded the configured threshold.
    ResidualTooHigh,
    /// Feature was never labelled by the detection pipeline (e.g. noise
    /// outside the recovered lattice support).
    Unlabelled,
    /// Feature was labelled by the topological pass but dropped by the
    /// post-build validation stage (line collinearity, local-H residual,
    /// or edge-length band).
    ValidationDropped,
}

/// Rejected feature record.
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub struct RejectedFeature {
    /// Caller-owned source index.
    pub source_index: usize,
    /// Coordinate associated with the rejection, if one was proposed.
    pub coord: Option<Coord>,
    /// Residual in image pixels, if available.
    pub residual_px: Option<f32>,
    /// Rejection reason.
    pub reason: RejectionReason,
}

impl RejectedFeature {
    /// Construct a rejected-feature record.
    pub fn new(
        source_index: usize,
        coord: Option<Coord>,
        residual_px: Option<f32>,
        reason: RejectionReason,
    ) -> Self {
        Self {
            source_index,
            coord,
            residual_px,
            reason,
        }
    }
}

/// Shared successful solution shape for grid tasks.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct GridSolution {
    /// Labelled grid entries.
    pub grid: LabelledGrid,
    /// Lattice fit, when the task computed one.
    pub fit: Option<LatticeFit>,
    /// Features rejected by task gates.
    pub rejected: Vec<RejectedFeature>,
}

impl GridSolution {
    /// Construct a grid solution.
    pub fn new(
        grid: LabelledGrid,
        fit: Option<LatticeFit>,
        rejected: Vec<RejectedFeature>,
    ) -> Self {
        Self {
            grid,
            fit,
            rejected,
        }
    }

    /// Linear-scan lookup of the rejection record for the given source index, if any.
    pub fn rejected_for(&self, source_index: usize) -> Option<&RejectedFeature> {
        self.rejected
            .iter()
            .find(|r| r.source_index == source_index)
    }
}

/// Report returned by coordinate-hypothesis consistency checks.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct ConsistencyReport {
    /// `true` when all residuals satisfy the configured threshold.
    pub passed: bool,
    /// Labelled solution and residual diagnostics.
    pub solution: GridSolution,
}

impl ConsistencyReport {
    /// Construct a consistency report.
    pub fn new(passed: bool, solution: GridSolution) -> Self {
        Self { passed, solution }
    }

    /// Convenience accessor for the maximum residual in pixels from the fitted lattice,
    /// when one was computed.
    pub fn max_residual_px(&self) -> Option<f32> {
        Some(self.solution.fit.as_ref()?.residuals.max_px)
    }
}

/// Shift every entry's coordinate so the bounding-box minimum is `(0, 0)`.
fn rebase_entries_to_origin(entries: &mut [GridEntry]) {
    if entries.is_empty() {
        return;
    }
    let (min_u, min_v) = entries.iter().fold((i32::MAX, i32::MAX), |(a, b), e| {
        (a.min(e.coord.u), b.min(e.coord.v))
    });
    if min_u != 0 || min_v != 0 {
        for e in entries.iter_mut() {
            e.coord.u -= min_u;
            e.coord.v -= min_v;
        }
    }
}

/// Permute / sign-flip the lattice axes so `+u` points roughly `+x` and `+v`
/// roughly `+y` in image pixels, keeping labels non-negative. Uses only
/// [`GridEntry::image_position`]; positions are unchanged.
///
/// The mean `+u` and `+v` step vectors are accumulated over adjacent labelled
/// pairs in a deterministic coordinate order (sorted keys), so the `f32` sums do
/// not depend on map iteration order — the swap / flip decision is a function of
/// signs and magnitude comparisons that is robust to ULP-level summation
/// differences.
/// Canonicalize entry labels to the image-axis frame, returning `true` when the
/// two lattice axes were transposed (`u` ↔ `v`). Callers that hold an axis-keyed
/// side value (e.g. [`GridDimensions`]) must apply the same transpose when this
/// returns `true`; sign flips alone (return `false`) keep the axis assignment.
fn canonicalize_to_image_axes(entries: &mut [GridEntry]) -> bool {
    if entries.len() < 2 {
        return false;
    }
    let pos_by_uv: HashMap<(i32, i32), (f32, f32)> = entries
        .iter()
        .map(|e| {
            (
                (e.coord.u, e.coord.v),
                (e.image_position.x, e.image_position.y),
            )
        })
        .collect();

    let mut keys: Vec<(i32, i32)> = pos_by_uv.keys().copied().collect();
    keys.sort_unstable();
    let mut vu_sum = (0.0_f32, 0.0_f32);
    let mut vv_sum = (0.0_f32, 0.0_f32);
    let mut vu_n = 0u32;
    let mut vv_n = 0u32;
    for &(u, v) in &keys {
        let (x, y) = pos_by_uv[&(u, v)];
        if let Some(&(xn, yn)) = pos_by_uv.get(&(u + 1, v)) {
            vu_sum.0 += xn - x;
            vu_sum.1 += yn - y;
            vu_n += 1;
        }
        if let Some(&(xn, yn)) = pos_by_uv.get(&(u, v + 1)) {
            vv_sum.0 += xn - x;
            vv_sum.1 += yn - y;
            vv_n += 1;
        }
    }
    if vu_n == 0 || vv_n == 0 {
        return false;
    }
    let vu = (vu_sum.0 / vu_n as f32, vu_sum.1 / vu_n as f32);
    let vv = (vv_sum.0 / vv_n as f32, vv_sum.1 / vv_n as f32);

    // Make the axis with the larger |x| component the horizontal (`u`) axis.
    let swap = vu.0.abs() < vv.0.abs();
    let new_vu = if swap { vv } else { vu };
    let new_vv = if swap { vu } else { vv };
    let flip_u = new_vu.0 < 0.0;
    let flip_v = new_vv.1 < 0.0;

    if !swap && !flip_u && !flip_v {
        return false;
    }

    // Post-swap extents, so the sign flip stays within the non-negative domain.
    let mut umax = i32::MIN;
    let mut vmax = i32::MIN;
    for e in entries.iter() {
        let (nu, nv) = if swap {
            (e.coord.v, e.coord.u)
        } else {
            (e.coord.u, e.coord.v)
        };
        umax = umax.max(nu);
        vmax = vmax.max(nv);
    }

    for e in entries.iter_mut() {
        let (mut nu, mut nv) = if swap {
            (e.coord.v, e.coord.u)
        } else {
            (e.coord.u, e.coord.v)
        };
        if flip_u {
            nu = umax - nu;
        }
        if flip_v {
            nv = vmax - nv;
        }
        e.coord.u = nu;
        e.coord.v = nv;
    }

    swap
}

fn bbox_for_entries(entries: &[GridEntry]) -> Option<(Coord, Coord)> {
    let first = entries.first()?;
    let mut min = first.coord;
    let mut max = first.coord;
    for entry in &entries[1..] {
        min.u = min.u.min(entry.coord.u);
        min.v = min.v.min(entry.coord.v);
        max.u = max.u.max(entry.coord.u);
        max.v = max.v.max(entry.coord.v);
    }
    Some((min, max))
}

#[cfg(test)]
mod tests {
    use nalgebra::{Point2, Projective2};

    use super::*;

    fn make_identity_fit() -> LatticeFit {
        LatticeFit::new(
            Projective2::identity(),
            ResidualSummary::new(1, 0.5_f32, 1.0_f32),
        )
    }

    #[test]
    fn max_residual_px_none_when_fit_absent() {
        let grid = LabelledGrid::new(LatticeKind::Square, vec![], None);
        let solution = GridSolution::new(grid, None, vec![]);
        let report = ConsistencyReport::new(true, solution);
        assert_eq!(report.max_residual_px(), None);
    }

    #[test]
    fn max_residual_px_some_when_fit_present() {
        let grid = LabelledGrid::new(LatticeKind::Square, vec![], None);
        let fit = make_identity_fit();
        let solution = GridSolution::new(grid, Some(fit), vec![]);
        let report = ConsistencyReport::new(true, solution);
        assert_eq!(report.max_residual_px(), Some(1.0_f32));
    }

    #[test]
    fn labelled_grid_find_present_and_absent() {
        let entry = GridEntry::new(Coord::new(0, 0), 42, Point2::new(1.0_f32, 2.0), None);
        let grid = LabelledGrid::new(LatticeKind::Square, vec![entry], None);
        assert!(grid.find(42).is_some());
        assert!(grid.find(99).is_none());
    }

    fn mk_entry(u: i32, v: i32, x: f32, y: f32) -> GridEntry {
        GridEntry::new(Coord::new(u, v), 0, Point2::new(x, y), None)
    }

    fn coord_by_pos(grid: &LabelledGrid) -> HashMap<(i32, i32), (i32, i32)> {
        grid.entries
            .iter()
            .map(|e| {
                (
                    (e.image_position.x as i32, e.image_position.y as i32),
                    (e.coord.u, e.coord.v),
                )
            })
            .collect()
    }

    #[test]
    fn normalize_rebases_and_sorts_already_canonical() {
        // +u already points +x, +v already points +y; only an offset to remove.
        let entries = vec![
            mk_entry(3, 5, 10.0, 10.0),
            mk_entry(4, 5, 20.0, 10.0),
            mk_entry(3, 6, 10.0, 20.0),
            mk_entry(4, 6, 20.0, 20.0),
        ];
        let mut grid = LabelledGrid::new(LatticeKind::Square, entries, None);
        grid.normalize();
        let by_pos = coord_by_pos(&grid);
        assert_eq!(by_pos[&(10, 10)], (0, 0));
        assert_eq!(by_pos[&(20, 10)], (1, 0));
        assert_eq!(by_pos[&(10, 20)], (0, 1));
        assert_eq!(by_pos[&(20, 20)], (1, 1));
        assert_eq!(grid.bbox, Some((Coord::new(0, 0), Coord::new(1, 1))));
        // Stable (v, u) order.
        let order: Vec<(i32, i32)> = grid
            .entries
            .iter()
            .map(|e| (e.coord.u, e.coord.v))
            .collect();
        assert_eq!(order, vec![(0, 0), (1, 0), (0, 1), (1, 1)]);
    }

    #[test]
    fn normalize_canonicalizes_rotated_axes() {
        // Builder assigned +u along +y and +v along +x (a 90° rotation);
        // normalize must swap so +u ≈ +x and +v ≈ +y, putting (0, 0) at the
        // smallest (x, y) corner.
        let entries = vec![
            mk_entry(0, 0, 10.0, 10.0),
            mk_entry(0, 1, 20.0, 10.0),
            mk_entry(1, 0, 10.0, 20.0),
            mk_entry(1, 1, 20.0, 20.0),
        ];
        let mut grid = LabelledGrid::new(LatticeKind::Square, entries, None);
        grid.normalize();
        let by_pos = coord_by_pos(&grid);
        assert_eq!(
            by_pos[&(10, 10)],
            (0, 0),
            "(0,0) must land at smallest (x,y)"
        );
        assert_eq!(by_pos[&(20, 10)], (1, 0), "+u must point +x");
        assert_eq!(by_pos[&(10, 20)], (0, 1), "+v must point +y");
    }

    #[test]
    fn normalize_transposes_dimensions_on_axis_swap() {
        // Same 90°-rotated geometry as the test above (builder put +u along +y
        // and +v along +x), so normalize swaps u ↔ v. Caller-supplied
        // rectangular dimensions must transpose with the axes, or they would
        // describe the pre-swap frame while entries/bbox describe the new one.
        let entries = vec![
            mk_entry(0, 0, 10.0, 10.0),
            mk_entry(0, 1, 20.0, 10.0),
            mk_entry(1, 0, 10.0, 20.0),
            mk_entry(1, 1, 20.0, 20.0),
        ];
        let mut grid = LabelledGrid::new(
            LatticeKind::Square,
            entries,
            Some(GridDimensions::new(5, 3)),
        );
        grid.normalize();
        assert_eq!(
            grid.dimensions,
            Some(GridDimensions::new(3, 5)),
            "axis swap must transpose width/height"
        );
    }

    #[test]
    fn normalize_keeps_dimensions_when_axes_only_flip() {
        // +u along -x, +v along +y: a sign flip on u, NO transpose. Dimensions
        // must be left untouched.
        let entries = vec![
            mk_entry(0, 0, 20.0, 10.0),
            mk_entry(1, 0, 10.0, 10.0),
            mk_entry(0, 1, 20.0, 20.0),
            mk_entry(1, 1, 10.0, 20.0),
        ];
        let mut grid = LabelledGrid::new(
            LatticeKind::Square,
            entries,
            Some(GridDimensions::new(5, 3)),
        );
        grid.normalize();
        assert_eq!(
            grid.dimensions,
            Some(GridDimensions::new(5, 3)),
            "a sign flip without a transpose must not touch dimensions"
        );
    }

    #[test]
    fn grid_solution_rejected_for_present_and_absent() {
        let rejected =
            RejectedFeature::new(5, None, Some(3.0_f32), RejectionReason::ResidualTooHigh);
        let grid = LabelledGrid::new(LatticeKind::Square, vec![], None);
        let solution = GridSolution::new(grid, None, vec![rejected]);
        assert!(solution.rejected_for(5).is_some());
        assert!(solution.rejected_for(0).is_none());
    }
}