// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! An indexed polygon signed-distance query, bit-identical to
//! [`polygon_signed_distance`].
//!
//! Measured motivation (2026-08-30 FSI3 profile): 51% of the fluid step
//! is `polygon_signed_distance` — the embedded mask rebuild and its
//! ghost reconstruction evaluate the SDF tens of thousands of times per
//! step, each walking every edge of a ~150-vertex interface polygon.
//! This index cuts each query to the handful of edges that can matter,
//! while returning EXACTLY the brute-force f64:
//!
//! - **Distance**: per-edge squared distances are computed by the same
//!   float ops as the brute force; the ring search visits a candidate
//!   set that provably contains the minimizing edge, and `min` over a
//!   superset containing the argmin equals `min` over all edges bit for
//!   bit. The ring lower bound uses the convex-projection inequality:
//!   for a query `q` clamped to `c` on the grid box and any point `z`
//!   inside it, `|q−z|² ≥ |q−c|² + |c−z|²`, so a ring at Chebyshev
//!   index `r` is at least `√(d_out² + ((r−1)·b)²)` away (`b` the
//!   smaller bin dimension).
//! - **Sign**: the even-odd ray test flips parity only for edges whose
//!   y-interval straddles the query, so edges are binned by y-interval
//!   and the query XORs over exactly the straddling candidates — the
//!   same tests, the same parity.
//!
//! The equality is asserted, not assumed: the tests compare against
//! [`polygon_signed_distance`] with `to_bits` over adversarial points.

use super::embedded_body::polygon_signed_distance;

/// Indexed signed distance to a closed polygon (negative inside, either
/// winding). Build once per geometry; queries are `O(edges near the
/// point)` instead of `O(all edges)`.
pub struct PolygonSdf {
    vertices: Vec<(f64, f64)>,
    /// Grid over the polygon's padded bounding box.
    x0: f64,
    y0: f64,
    bin_w: f64,
    bin_h: f64,
    nx: usize,
    ny: usize,
    /// Edge indices per bin (an edge appears in every bin its padded
    /// bounding box overlaps).
    bins: Vec<Vec<u32>>,
    /// Edge indices per y-row of the SAME grid, for the ray-crossing
    /// parity: an edge appears in every row its y-interval overlaps.
    rows: Vec<Vec<u32>>,
}

impl PolygonSdf {
    /// Index `vertices` (at least 3). Bin count scales with the edge
    /// count so construction stays `O(edges)`.
    #[must_use]
    pub fn new(vertices: Vec<(f64, f64)>) -> Self {
        assert!(vertices.len() >= 3, "a polygon needs at least 3 vertices");
        let n = vertices.len();
        let (mut min_x, mut min_y) = (f64::MAX, f64::MAX);
        let (mut max_x, mut max_y) = (f64::MIN, f64::MIN);
        for &(x, y) in &vertices {
            min_x = min_x.min(x);
            min_y = min_y.min(y);
            max_x = max_x.max(x);
            max_y = max_y.max(y);
        }
        // Degenerate extents still get a positive bin size.
        let width = (max_x - min_x).max(1e-12);
        let height = (max_y - min_y).max(1e-12);
        // ~2 edges per bin on a perimeter polygon: bins ~ n along the
        // longer side, aspect-scaled on the shorter.
        let nx = ((n as f64).sqrt() * (width / height).sqrt().max(0.25))
            .ceil()
            .clamp(1.0, 256.0) as usize;
        let ny = ((n as f64).sqrt() * (height / width).sqrt().max(0.25))
            .ceil()
            .clamp(1.0, 256.0) as usize;
        let bin_w = width / nx as f64;
        let bin_h = height / ny as f64;

        let mut bins = vec![Vec::new(); nx * ny];
        let mut rows = vec![Vec::new(); ny];
        let clamp_i = |x: f64| (((x - min_x) / bin_w) as isize).clamp(0, nx as isize - 1) as usize;
        let clamp_j = |y: f64| (((y - min_y) / bin_h) as isize).clamp(0, ny as isize - 1) as usize;
        for k in 0..n {
            let (ax, ay) = vertices[k];
            let (bx, by) = vertices[(k + 1) % n];
            let (i0, i1) = (clamp_i(ax.min(bx)), clamp_i(ax.max(bx)));
            let (j0, j1) = (clamp_j(ay.min(by)), clamp_j(ay.max(by)));
            for j in j0..=j1 {
                for i in i0..=i1 {
                    bins[j * nx + i].push(k as u32);
                }
                rows[j].push(k as u32);
            }
        }
        Self {
            vertices,
            x0: min_x,
            y0: min_y,
            bin_w,
            bin_h,
            nx,
            ny,
            bins,
            rows,
        }
    }

    /// The indexed vertices.
    #[must_use]
    pub fn vertices(&self) -> &[(f64, f64)] {
        &self.vertices
    }

    /// Squared distance from `(x, y)` to edge `k` — float-op for
    /// float-op the brute force's per-edge computation.
    #[inline]
    fn edge_dist2(&self, k: u32, x: f64, y: f64) -> f64 {
        let n = self.vertices.len();
        let (ax, ay) = self.vertices[k as usize];
        let (bx, by) = self.vertices[(k as usize + 1) % n];
        let (ex, ey) = (bx - ax, by - ay);
        let len2 = ex * ex + ey * ey;
        let s = if len2 > 0.0 {
            (((x - ax) * ex + (y - ay) * ey) / len2).clamp(0.0, 1.0)
        } else {
            0.0
        };
        let (qx, qy) = (ax + s * ex - x, ay + s * ey - y);
        qx * qx + qy * qy
    }

    /// Signed distance, bit-identical to
    /// `polygon_signed_distance(self.vertices(), x, y)`.
    #[must_use]
    pub fn signed_distance(&self, x: f64, y: f64) -> f64 {
        // Sign: XOR the ray test over the y-row candidates. Any edge
        // that straddles y lies in this row's list (its y-interval
        // overlaps the row), so the parity is over exactly the edges
        // the brute force flips on.
        let j_row = (((y - self.y0) / self.bin_h) as isize).clamp(0, self.ny as isize - 1) as usize;
        let mut inside = false;
        // Edges whose y-interval leaves the grid entirely are impossible
        // (the grid spans the polygon's bbox), but a query y outside the
        // bbox straddles nothing — the clamped row still contains every
        // straddling edge because there are none.
        for &k in &self.rows[j_row] {
            let n = self.vertices.len();
            let (ax, ay) = self.vertices[k as usize];
            let (bx, by) = self.vertices[(k as usize + 1) % n];
            if (ay > y) != (by > y) {
                let x_cross = ax + (y - ay) / (by - ay) * (bx - ax);
                if x < x_cross {
                    inside = !inside;
                }
            }
        }

        // Distance: ring search from the clamped bin.
        let ci = (((x - self.x0) / self.bin_w) as isize).clamp(0, self.nx as isize - 1);
        let cj = (((y - self.y0) / self.bin_h) as isize).clamp(0, self.ny as isize - 1);
        // Distance from the query to the grid box (0 inside).
        let cx = x.clamp(self.x0, self.x0 + self.bin_w * self.nx as f64);
        let cy = y.clamp(self.y0, self.y0 + self.bin_h * self.ny as f64);
        let d_out2 = (x - cx) * (x - cx) + (y - cy) * (y - cy);
        let b = self.bin_w.min(self.bin_h);

        let mut dist2 = f64::MAX;
        let max_ring = self.nx.max(self.ny) as isize;
        for r in 0..=max_ring {
            // Every point of a ring-r bin is at least this far away
            // (convex-projection inequality; see the module docs).
            if r >= 2 {
                let lb = (r - 1) as f64 * b;
                if d_out2 + lb * lb > dist2 {
                    break;
                }
            }
            let (i_lo, i_hi) = (ci - r, ci + r);
            let (j_lo, j_hi) = (cj - r, cj + r);
            let mut visit = |i: isize, j: isize, dist2: &mut f64| {
                if i < 0 || j < 0 || i >= self.nx as isize || j >= self.ny as isize {
                    return;
                }
                for &k in &self.bins[j as usize * self.nx + i as usize] {
                    let d2 = self.edge_dist2(k, x, y);
                    if d2 < *dist2 {
                        *dist2 = d2;
                    }
                }
            };
            if r == 0 {
                visit(ci, cj, &mut dist2);
            } else {
                for i in i_lo..=i_hi {
                    visit(i, j_lo, &mut dist2);
                    visit(i, j_hi, &mut dist2);
                }
                for j in (j_lo + 1)..j_hi {
                    visit(i_lo, j, &mut dist2);
                    visit(i_hi, j, &mut dist2);
                }
            }
        }

        let dist = dist2.sqrt();
        if inside { -dist } else { dist }
    }
}

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

    /// A deterministic pseudo-random stream (no rand dependency).
    struct Lcg(u64);
    impl Lcg {
        fn next_f64(&mut self, lo: f64, hi: f64) -> f64 {
            self.0 = self
                .0
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            let u = (self.0 >> 11) as f64 / (1u64 << 53) as f64;
            lo + u * (hi - lo)
        }
    }

    fn assert_bit_identical(vertices: &[(f64, f64)], points: &[(f64, f64)]) {
        let sdf = PolygonSdf::new(vertices.to_vec());
        for &(x, y) in points {
            let brute = polygon_signed_distance(vertices, x, y);
            let indexed = sdf.signed_distance(x, y);
            assert_eq!(
                brute.to_bits(),
                indexed.to_bits(),
                "indexed {indexed:.17e} != brute {brute:.17e} at ({x}, {y})"
            );
        }
    }

    /// A flag-like polygon: a long thin rectangle sampled densely (the
    /// FSI interface walk's shape), mildly deformed.
    fn flag_polygon(n_per_side: usize, deflect: f64) -> Vec<(f64, f64)> {
        let (x0, x1, y0, y1) = (0.25, 0.6, 0.19, 0.21);
        let mut v = Vec::new();
        for k in 0..n_per_side {
            let s = k as f64 / n_per_side as f64;
            let x = x0 + s * (x1 - x0);
            v.push((x, y0 + deflect * s * s));
        }
        for k in 0..3 {
            let s = k as f64 / 3.0;
            v.push((x1, y0 + deflect + s * (y1 - y0)));
        }
        for k in 0..n_per_side {
            let s = k as f64 / n_per_side as f64;
            let x = x1 - s * (x1 - x0);
            v.push((x, y1 + deflect * (1.0 - s) * (1.0 - s)));
        }
        v.push((x0, y1));
        v
    }

    #[test]
    fn bit_identical_on_flag_polygon() {
        for &deflect in &[0.0, 0.05, -0.08] {
            let vertices = flag_polygon(72, deflect);
            let mut rng = Lcg(42);
            let mut points = Vec::new();
            // The whole domain, the near field, and exactly-on-feature
            // points (vertices, edge midpoints, the bbox corners).
            for _ in 0..2000 {
                points.push((rng.next_f64(0.0, 2.5), rng.next_f64(0.0, 0.41)));
            }
            for _ in 0..2000 {
                points.push((rng.next_f64(0.24, 0.62), rng.next_f64(0.15, 0.28)));
            }
            for k in 0..vertices.len() {
                let (ax, ay) = vertices[k];
                let (bx, by) = vertices[(k + 1) % vertices.len()];
                points.push((ax, ay));
                points.push((0.5 * (ax + bx), 0.5 * (ay + by)));
            }
            points.push((-3.0, -1.0));
            points.push((10.0, 5.0));
            assert_bit_identical(&vertices, &points);
        }
    }

    #[test]
    fn bit_identical_on_random_polygons() {
        let mut rng = Lcg(7);
        for poly in 0..20 {
            let n = 3 + (poly % 9);
            let mut vertices: Vec<(f64, f64)> = (0..n)
                .map(|_| (rng.next_f64(-1.0, 1.0), rng.next_f64(-1.0, 1.0)))
                .collect();
            // Exercise degenerate zero-length edges too.
            if poly % 4 == 0 {
                let first = vertices[0];
                vertices.insert(1, first);
            }
            let points: Vec<(f64, f64)> = (0..1500)
                .map(|_| (rng.next_f64(-3.0, 3.0), rng.next_f64(-3.0, 3.0)))
                .collect();
            assert_bit_identical(&vertices, &points);
        }
    }

    #[test]
    fn bit_identical_on_horizontal_edge_rays() {
        // Horizontal edges never straddle their own y (the strict/loose
        // comparison pair `(ay > y) != (by > y)` is false when ay == by),
        // and queries exactly AT a vertex y exercise the boundary of the
        // straddle test. The parity must match the brute force on all of
        // them.
        let vertices = vec![
            (0.0, 0.0),
            (2.0, 0.0),
            (2.0, 1.0),
            (1.0, 1.0),
            (1.0, 0.5),
            (0.0, 0.5),
        ];
        let mut points = Vec::new();
        for &y in &[0.0, 0.25, 0.5, 0.75, 1.0] {
            for k in 0..40 {
                points.push((-0.5 + 3.0 * k as f64 / 39.0, y));
            }
        }
        assert_bit_identical(&vertices, &points);
    }
}