projective-grid 0.4.2

//! Per-triangle homography mesh for hex grid rectification.
//!
//! Given a map of hex grid corners (axial coordinates) to image positions,
//! builds one affine transform and one homography per triangle cell.
//! The hex lattice is decomposed into parallelogram cells, each split into
//! two triangles.

use crate::grid_index::GridIndex;
use crate::homography::{estimate_homography, Homography};
use nalgebra::{Matrix2, Point2, Vector2};
use std::collections::HashMap;

/// Sqrt(3) / 2, the vertical spacing factor for pointy-top hex grids.
const SQRT3_HALF: f64 = 0.866_025_403_784_438_6;

#[derive(thiserror::Error, Debug)]
pub enum HexMeshError {
    #[error("not enough grid corners (need at least 3)")]
    NotEnoughCorners,
    #[error("no valid triangles found")]
    NoValidTriangles,
}

/// A 2D affine transform: `dst = M * [src_x, src_y]^T + t`.
#[derive(Clone, Copy, Debug)]
pub struct AffineTransform2D {
    /// 2x2 linear part.
    pub linear: Matrix2<f64>,
    /// Translation part.
    pub translation: Vector2<f64>,
}

impl AffineTransform2D {
    /// Compute the affine transform mapping `src` triangle to `dst` triangle.
    ///
    /// Returns `None` if the source triangle is degenerate (collinear points).
    pub fn from_triangle_correspondence(
        src: [Point2<f64>; 3],
        dst: [Point2<f64>; 3],
    ) -> Option<Self> {
        // Solve: dst_i = M * src_i + t for i = 0, 1, 2
        // Using src[0] as origin: M * (src_i - src_0) = (dst_i - dst_0) for i = 1, 2
        let ds1 = src[1] - src[0];
        let ds2 = src[2] - src[0];
        let dd1 = dst[1] - dst[0];
        let dd2 = dst[2] - dst[0];

        // [ds1 | ds2] as column matrix, invert to get M
        let src_mat = Matrix2::new(ds1.x, ds2.x, ds1.y, ds2.y);

        let src_inv = src_mat.try_inverse()?;

        // M = [dd1 | dd2] * src_inv
        let dst_mat = Matrix2::new(dd1.x, dd2.x, dd1.y, dd2.y);
        let linear = dst_mat * src_inv;

        let t = dst[0] - linear * Vector2::new(src[0].x, src[0].y);
        let translation = Vector2::new(t.x, t.y);

        Some(Self {
            linear,
            translation,
        })
    }

    /// Apply the transform to a 2D point.
    pub fn apply(&self, p: Point2<f64>) -> Point2<f64> {
        let v = self.linear * Vector2::new(p.x, p.y) + self.translation;
        Point2::new(v.x, v.y)
    }
}

#[derive(Clone, Debug)]
struct TriangleCell {
    /// Affine transform from rectified triangle to image triangle.
    affine: AffineTransform2D,
    /// Homography from rectified triangle to image triangle (4-point, with centroid).
    homography: Homography,
}

/// Per-triangle homography mesh over a hex grid.
///
/// Each parallelogram cell in axial space `(q, r) → (q+1, r+1)` is split
/// into two triangles:
/// - **Lower**: `(q,r)`, `(q+1,r)`, `(q,r+1)` — when `frac_q + frac_r ≤ 1`
/// - **Upper**: `(q+1,r)`, `(q,r+1)`, `(q+1,r+1)` — when `frac_q + frac_r > 1`
#[derive(Clone, Debug)]
pub struct HexGridHomographyMesh {
    pub min_q: i32,
    pub min_r: i32,
    /// Number of parallelogram cells along q.
    pub cells_q: usize,
    /// Number of parallelogram cells along r.
    pub cells_r: usize,
    /// Rectified pixels per grid cell edge.
    pub px_per_cell: f32,
    /// Number of valid triangle cells.
    pub valid_triangles: usize,
    /// Rectified image dimensions.
    pub rect_width: usize,
    pub rect_height: usize,

    // 2 triangles per parallelogram cell: [lower, upper] interleaved.
    // Length = cells_q * cells_r * 2
    cells: Vec<Option<TriangleCell>>,

    // Rectified coordinate offset (subtracted from raw axial→rect mapping).
    x_offset: f64,
    y_offset: f64,
}

impl HexGridHomographyMesh {
    /// Build per-triangle transforms from a hex grid corner map.
    ///
    /// - `corners`: map from axial grid index `(q=i, r=j)` to image position.
    /// - `px_per_cell`: rectified pixels per grid cell edge.
    pub fn from_corners(
        corners: &HashMap<GridIndex, Point2<f32>>,
        px_per_cell: f32,
    ) -> Result<Self, HexMeshError> {
        if corners.len() < 3 {
            return Err(HexMeshError::NotEnoughCorners);
        }

        let (mut min_q, mut min_r) = (i32::MAX, i32::MAX);
        let (mut max_q, mut max_r) = (i32::MIN, i32::MIN);
        for g in corners.keys() {
            min_q = min_q.min(g.i);
            min_r = min_r.min(g.j);
            max_q = max_q.max(g.i);
            max_r = max_r.max(g.j);
        }

        if max_q - min_q < 1 || max_r - min_r < 1 {
            return Err(HexMeshError::NoValidTriangles);
        }

        let cells_q = (max_q - min_q) as usize;
        let cells_r = (max_r - min_r) as usize;
        let s = px_per_cell as f64;

        // Compute rectified bounding box
        let mut x_min = f64::MAX;
        let mut x_max = f64::MIN;
        let mut y_min = f64::MAX;
        let mut y_max = f64::MIN;

        // Check all corner positions of the bounding parallelogram
        for &q in &[min_q, max_q] {
            for &r in &[min_r, max_r] {
                let x = s * (q as f64 + r as f64 * 0.5);
                let y = s * (r as f64 * SQRT3_HALF);
                x_min = x_min.min(x);
                x_max = x_max.max(x);
                y_min = y_min.min(y);
                y_max = y_max.max(y);
            }
        }

        let rect_width = ((x_max - x_min).round().max(1.0)) as usize;
        let rect_height = ((y_max - y_min).round().max(1.0)) as usize;

        let axial_to_rect = |q: i32, r: i32| -> Point2<f64> {
            Point2::new(
                s * (q as f64 + r as f64 * 0.5) - x_min,
                s * (r as f64 * SQRT3_HALF) - y_min,
            )
        };

        let mut cells = vec![None; cells_q * cells_r * 2];
        let mut valid_triangles = 0usize;

        for cr in 0..cells_r {
            for cq in 0..cells_q {
                let q0 = min_q + cq as i32;
                let r0 = min_r + cr as i32;

                let g00 = GridIndex { i: q0, j: r0 };
                let g10 = GridIndex { i: q0 + 1, j: r0 };
                let g01 = GridIndex { i: q0, j: r0 + 1 };
                let g11 = GridIndex {
                    i: q0 + 1,
                    j: r0 + 1,
                };

                let p00 = corners.get(&g00).copied();
                let p10 = corners.get(&g10).copied();
                let p01 = corners.get(&g01).copied();
                let p11 = corners.get(&g11).copied();

                let idx_base = (cr * cells_q + cq) * 2;

                // Lower triangle: g00, g10, g01
                if let (Some(ip00), Some(ip10), Some(ip01)) = (p00, p10, p01) {
                    let rect_tri = [
                        axial_to_rect(q0, r0),
                        axial_to_rect(q0 + 1, r0),
                        axial_to_rect(q0, r0 + 1),
                    ];
                    let img_tri = [
                        Point2::new(ip00.x as f64, ip00.y as f64),
                        Point2::new(ip10.x as f64, ip10.y as f64),
                        Point2::new(ip01.x as f64, ip01.y as f64),
                    ];

                    if let Some(affine) =
                        AffineTransform2D::from_triangle_correspondence(rect_tri, img_tri)
                    {
                        // 4-point homography: add centroid as 4th point
                        let rect_c = centroid(&rect_tri);
                        let img_c = affine.apply(rect_c);
                        let rect_4: Vec<Point2<f32>> = rect_tri
                            .iter()
                            .chain(std::iter::once(&rect_c))
                            .map(|p| Point2::new(p.x as f32, p.y as f32))
                            .collect();
                        let img_4: Vec<Point2<f32>> = img_tri
                            .iter()
                            .chain(std::iter::once(&img_c))
                            .map(|p| Point2::new(p.x as f32, p.y as f32))
                            .collect();

                        if let Some(homography) = estimate_homography(&rect_4, &img_4) {
                            cells[idx_base] = Some(TriangleCell { affine, homography });
                            valid_triangles += 1;
                        }
                    }
                }

                // Upper triangle: g10, g01, g11
                if let (Some(ip10), Some(ip01), Some(ip11)) = (p10, p01, p11) {
                    let rect_tri = [
                        axial_to_rect(q0 + 1, r0),
                        axial_to_rect(q0, r0 + 1),
                        axial_to_rect(q0 + 1, r0 + 1),
                    ];
                    let img_tri = [
                        Point2::new(ip10.x as f64, ip10.y as f64),
                        Point2::new(ip01.x as f64, ip01.y as f64),
                        Point2::new(ip11.x as f64, ip11.y as f64),
                    ];

                    if let Some(affine) =
                        AffineTransform2D::from_triangle_correspondence(rect_tri, img_tri)
                    {
                        let rect_c = centroid(&rect_tri);
                        let img_c = affine.apply(rect_c);
                        let rect_4: Vec<Point2<f32>> = rect_tri
                            .iter()
                            .chain(std::iter::once(&rect_c))
                            .map(|p| Point2::new(p.x as f32, p.y as f32))
                            .collect();
                        let img_4: Vec<Point2<f32>> = img_tri
                            .iter()
                            .chain(std::iter::once(&img_c))
                            .map(|p| Point2::new(p.x as f32, p.y as f32))
                            .collect();

                        if let Some(homography) = estimate_homography(&rect_4, &img_4) {
                            cells[idx_base + 1] = Some(TriangleCell { affine, homography });
                            valid_triangles += 1;
                        }
                    }
                }
            }
        }

        if valid_triangles == 0 {
            return Err(HexMeshError::NoValidTriangles);
        }

        Ok(Self {
            min_q,
            min_r,
            cells_q,
            cells_r,
            px_per_cell,
            valid_triangles,
            rect_width,
            rect_height,
            cells,
            x_offset: x_min,
            y_offset: y_min,
        })
    }

    /// Map a point in **global rectified pixel coordinates** to image coordinates
    /// using the per-triangle affine transform.
    ///
    /// Returns `None` if the point lies outside the mesh or the cell is invalid.
    pub fn rect_to_img_affine(&self, p_rect: Point2<f32>) -> Option<Point2<f32>> {
        let (cell, p64) = self.lookup_cell(p_rect)?;
        let result = cell.affine.apply(p64);
        Some(Point2::new(result.x as f32, result.y as f32))
    }

    /// Map a point in **global rectified pixel coordinates** to image coordinates
    /// using the per-triangle homography.
    ///
    /// Returns `None` if the point lies outside the mesh or the cell is invalid.
    pub fn rect_to_img(&self, p_rect: Point2<f32>) -> Option<Point2<f32>> {
        let (cell, _) = self.lookup_cell(p_rect)?;
        Some(cell.homography.apply(p_rect))
    }

    /// Look up the triangle cell for a rectified point.
    fn lookup_cell(&self, p_rect: Point2<f32>) -> Option<(&TriangleCell, Point2<f64>)> {
        let s = self.px_per_cell as f64;
        if s <= 0.0 {
            return None;
        }

        let p64 = Point2::new(p_rect.x as f64, p_rect.y as f64);

        // Convert rectified pixel coords back to fractional axial coords
        let r_frac = (p64.y + self.y_offset) / (s * SQRT3_HALF);
        let q_frac = (p64.x + self.x_offset) / s - r_frac * 0.5;

        // Determine parallelogram cell
        let cq_f = q_frac - self.min_q as f64;
        let cr_f = r_frac - self.min_r as f64;

        let cq = cq_f.floor() as i32;
        let cr = cr_f.floor() as i32;

        if cq < 0 || cr < 0 || cq >= self.cells_q as i32 || cr >= self.cells_r as i32 {
            return None;
        }

        // Determine lower vs upper triangle
        let frac_q = cq_f - cq as f64;
        let frac_r = cr_f - cr as f64;
        let is_upper = frac_q + frac_r > 1.0;

        let idx = (cr as usize * self.cells_q + cq as usize) * 2 + is_upper as usize;
        let cell = self.cells.get(idx)?.as_ref()?;

        Some((cell, p64))
    }
}

fn centroid(tri: &[Point2<f64>; 3]) -> Point2<f64> {
    Point2::new(
        (tri[0].x + tri[1].x + tri[2].x) / 3.0,
        (tri[0].y + tri[1].y + tri[2].y) / 3.0,
    )
}

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

    fn make_hex_corners(radius: i32, spacing: f32) -> HashMap<GridIndex, Point2<f32>> {
        let sqrt3 = 3.0f32.sqrt();
        let mut map = HashMap::new();
        for q in -radius..=radius {
            for r in -radius..=radius {
                if (q + r).abs() > radius {
                    continue;
                }
                let x = spacing * (q as f32 + r as f32 * 0.5);
                let y = spacing * (r as f32 * sqrt3 / 2.0);
                map.insert(GridIndex { i: q, j: r }, Point2::new(x, y));
            }
        }
        map
    }

    #[test]
    fn affine_from_triangle_identity() {
        let tri = [
            Point2::new(0.0, 0.0),
            Point2::new(1.0, 0.0),
            Point2::new(0.0, 1.0),
        ];
        let aff = AffineTransform2D::from_triangle_correspondence(tri, tri).unwrap();
        let p = Point2::new(0.3, 0.4);
        let result = aff.apply(p);
        assert!((result.x - p.x).abs() < 1e-10);
        assert!((result.y - p.y).abs() < 1e-10);
    }

    #[test]
    fn affine_maps_vertices_correctly() {
        let src = [
            Point2::new(0.0, 0.0),
            Point2::new(1.0, 0.0),
            Point2::new(0.0, 1.0),
        ];
        let dst = [
            Point2::new(10.0, 20.0),
            Point2::new(30.0, 20.0),
            Point2::new(10.0, 50.0),
        ];
        let aff = AffineTransform2D::from_triangle_correspondence(src, dst).unwrap();
        for (s, d) in src.iter().zip(dst.iter()) {
            let result = aff.apply(*s);
            assert!((result.x - d.x).abs() < 1e-10);
            assert!((result.y - d.y).abs() < 1e-10);
        }
    }

    #[test]
    fn degenerate_triangle_returns_none() {
        let src = [
            Point2::new(0.0, 0.0),
            Point2::new(1.0, 0.0),
            Point2::new(2.0, 0.0), // collinear
        ];
        let dst = src;
        assert!(AffineTransform2D::from_triangle_correspondence(src, dst).is_none());
    }

    #[test]
    fn mesh_from_regular_hex_grid() {
        let corners = make_hex_corners(3, 60.0);
        let mesh = HexGridHomographyMesh::from_corners(&corners, 60.0).unwrap();
        assert!(mesh.valid_triangles > 0);
        assert!(mesh.rect_width > 0);
        assert!(mesh.rect_height > 0);
    }

    #[test]
    fn round_trip_through_affine_mesh() {
        let spacing = 60.0;
        let corners = make_hex_corners(3, spacing);
        let mesh = HexGridHomographyMesh::from_corners(&corners, spacing).unwrap();

        // Test that known corner positions round-trip through the mesh
        let s = spacing as f64;

        // Verify that corners at known positions map back reasonably
        for (g, &img_pos) in &corners {
            let rx = (s * (g.i as f64 + g.j as f64 * 0.5) - mesh.x_offset) as f32;
            let ry = (s * (g.j as f64 * SQRT3_HALF) - mesh.y_offset) as f32;
            let rect_pt = Point2::new(rx, ry);

            if let Some(recovered) = mesh.rect_to_img_affine(rect_pt) {
                assert!(
                    (recovered.x - img_pos.x).abs() < 1.0,
                    "x mismatch at ({},{}): {} vs {}",
                    g.i,
                    g.j,
                    recovered.x,
                    img_pos.x,
                );
                assert!(
                    (recovered.y - img_pos.y).abs() < 1.0,
                    "y mismatch at ({},{}): {} vs {}",
                    g.i,
                    g.j,
                    recovered.y,
                    img_pos.y,
                );
            }
            // Some boundary corners may not have a valid triangle cell — that's OK
        }
    }

    #[test]
    fn round_trip_through_homography_mesh() {
        let spacing = 60.0;
        let corners = make_hex_corners(3, spacing);
        let mesh = HexGridHomographyMesh::from_corners(&corners, spacing).unwrap();

        let s = spacing as f64;

        for (g, &img_pos) in &corners {
            let rx = (s * (g.i as f64 + g.j as f64 * 0.5) - mesh.x_offset) as f32;
            let ry = (s * (g.j as f64 * SQRT3_HALF) - mesh.y_offset) as f32;
            let rect_pt = Point2::new(rx, ry);

            if let Some(recovered) = mesh.rect_to_img(rect_pt) {
                assert!(
                    (recovered.x - img_pos.x).abs() < 1.0,
                    "homography x mismatch at ({},{}): {} vs {}",
                    g.i,
                    g.j,
                    recovered.x,
                    img_pos.x,
                );
                assert!(
                    (recovered.y - img_pos.y).abs() < 1.0,
                    "homography y mismatch at ({},{}): {} vs {}",
                    g.i,
                    g.j,
                    recovered.y,
                    img_pos.y,
                );
            }
        }
    }

    #[test]
    fn too_few_corners_errors() {
        let mut corners = HashMap::new();
        corners.insert(GridIndex { i: 0, j: 0 }, Point2::new(0.0, 0.0));
        corners.insert(GridIndex { i: 1, j: 0 }, Point2::new(50.0, 0.0));

        let result = HexGridHomographyMesh::from_corners(&corners, 50.0);
        assert!(result.is_err());
    }

    #[test]
    fn missing_corners_handled_gracefully() {
        let mut corners = make_hex_corners(3, 60.0);
        // Remove some corners
        corners.remove(&GridIndex { i: 0, j: 0 });
        corners.remove(&GridIndex { i: 1, j: 1 });

        let mesh = HexGridHomographyMesh::from_corners(&corners, 60.0);
        // Should still succeed (just with fewer valid triangles)
        assert!(mesh.is_ok());
        let mesh = mesh.unwrap();
        assert!(mesh.valid_triangles > 0);
    }
}