tetra3 0.5.1

//! Polynomial (SIP-like) distortion model.
//!
//! Models arbitrary 2D distortion using polynomial correction terms:
//!
//! ```text
//! x_distorted = x + Σ A_pq · x^p · y^q     (0 ≤ p+q ≤ order)
//! y_distorted = y + Σ B_pq · x^p · y^q
//! ```
//!
//! Including all terms from order 0:
//! - **(p+q = 0)**: constant offset — optical center shift
//! - **(p+q = 1)**: linear terms  — residual scale & rotation
//! - **(p+q ≥ 2)**: higher-order distortion
//!
//! Unlike the radial model, this captures tangential distortion, decentering,
//! and other effects that aren't radially symmetric — critical for cameras
//! like TESS where each CCD is offset from the optical axis.
//!
//! Inverse distortion uses Newton iteration on the forward polynomial — see
//! [`PolynomialDistortion::undistort`]. The legacy `ap_coeffs` / `bp_coeffs`
//! fields remain in the struct for binary-format compatibility but are
//! zero-valued in any model produced by this crate.
//!
//! Coordinates are in pixels relative to the image center. The coefficients
//! are stored normalized: internally the (x, y) inputs are divided by a
//! `scale` factor (typically half the image width) before evaluating the
//! polynomial, so the coefficients stay in a numerically well-conditioned range.

/// SIP-like polynomial distortion model.
///
/// Forward distortion (ideal → distorted) is the explicit polynomial.
/// Inverse distortion (distorted → ideal) is computed by Newton iteration
/// on the forward polynomial — see [`Self::undistort`].
#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
pub struct PolynomialDistortion {
    /// Polynomial order (2..=6 typically).
    pub order: u32,
    /// Normalization scale: coordinates are divided by this before evaluation.
    /// Typically image_width / 2.
    pub scale: f64,
    /// Forward A coefficients (x correction, ideal → distorted) in normalized coords.
    /// Stored as a flat vector; use `coeff_index(p, q)` to access.
    pub a_coeffs: Vec<f64>,
    /// Forward B coefficients (y correction, ideal → distorted) in normalized coords.
    pub b_coeffs: Vec<f64>,
    /// Legacy inverse AP coefficients. Zero in models produced by this crate;
    /// retained in the struct only for binary-format compatibility with
    /// previously-saved camera models.
    pub ap_coeffs: Vec<f64>,
    /// Legacy inverse BP coefficients. See [`Self::ap_coeffs`].
    pub bp_coeffs: Vec<f64>,
}

impl PolynomialDistortion {
    /// Create a new polynomial distortion model.
    ///
    /// All coefficient vectors must have exactly `num_coeffs(order)` elements.
    pub fn new(
        order: u32,
        scale: f64,
        a_coeffs: Vec<f64>,
        b_coeffs: Vec<f64>,
        ap_coeffs: Vec<f64>,
        bp_coeffs: Vec<f64>,
    ) -> Self {
        let n = num_coeffs(order);
        assert_eq!(a_coeffs.len(), n, "a_coeffs length mismatch");
        assert_eq!(b_coeffs.len(), n, "b_coeffs length mismatch");
        assert_eq!(ap_coeffs.len(), n, "ap_coeffs length mismatch");
        assert_eq!(bp_coeffs.len(), n, "bp_coeffs length mismatch");
        Self {
            order,
            scale,
            a_coeffs,
            b_coeffs,
            ap_coeffs,
            bp_coeffs,
        }
    }

    /// Create a zero (identity) polynomial distortion model.
    pub fn zero(order: u32, scale: f64) -> Self {
        let n = num_coeffs(order);
        Self {
            order,
            scale,
            a_coeffs: vec![0.0; n],
            b_coeffs: vec![0.0; n],
            ap_coeffs: vec![0.0; n],
            bp_coeffs: vec![0.0; n],
        }
    }

    /// Forward distortion: ideal → distorted (pixel coords, relative to image center).
    pub fn distort(&self, x: f64, y: f64) -> (f64, f64) {
        let u = x / self.scale;
        let v = y / self.scale;
        let dx = eval_poly(&self.a_coeffs, self.order, u, v);
        let dy = eval_poly(&self.b_coeffs, self.order, u, v);
        (x + dx * self.scale, y + dy * self.scale)
    }

    /// Inverse distortion: distorted → ideal (pixel coords, relative to image center).
    ///
    /// Uses Newton iteration on the **forward** polynomial. Given an observed
    /// distorted pixel `(x_d, y_d)`, finds the ideal `(x, y)` such that
    /// `distort(x, y) = (x_d, y_d)`. Converges in 2–4 iterations to machine
    /// precision for typical lens distortion (sub-pixel correction terms).
    ///
    /// This is exact (limited only by the forward polynomial's expressiveness),
    /// in contrast to evaluating a separately-fit inverse polynomial which has
    /// a small "asymmetry" error because a finite-order polynomial cannot
    /// perfectly invert another finite-order polynomial.
    pub fn undistort(&self, x_d: f64, y_d: f64) -> (f64, f64) {
        const MAX_ITER: usize = 8;
        const TOL_PX: f64 = 1e-9; // sub-nanopixel residual

        // Initial guess: pixels are close enough to ideal that x_d ≈ x.
        // (For TESS the worst-case distortion is ~10 px on 2048-wide images.)
        let mut x = x_d;
        let mut y = y_d;

        for _ in 0..MAX_ITER {
            let u = x / self.scale;
            let v = y / self.scale;
            let (a_val, da_du, da_dv) = eval_poly_with_grad(&self.a_coeffs, self.order, u, v);
            let (b_val, db_du, db_dv) = eval_poly_with_grad(&self.b_coeffs, self.order, u, v);

            // Forward distortion: F(x, y) = (x + s·A(u, v), y + s·B(u, v))
            let fx = x + a_val * self.scale;
            let fy = y + b_val * self.scale;

            // Residual: F(x, y) − (x_d, y_d)
            let rx = fx - x_d;
            let ry = fy - y_d;
            if rx * rx + ry * ry < TOL_PX * TOL_PX {
                break;
            }

            // Jacobian: ∂F/∂(x, y).
            //   ∂(s·A(x/s, y/s))/∂x = s · (1/s) · ∂A/∂u = ∂A/∂u
            // So J = [[1 + ∂A/∂u, ∂A/∂v], [∂B/∂u, 1 + ∂B/∂v]].
            let j11 = 1.0 + da_du;
            let j12 = da_dv;
            let j21 = db_du;
            let j22 = 1.0 + db_dv;
            let det = j11 * j22 - j12 * j21;
            // Singular Jacobian indicates near-degenerate distortion at this
            // point. We've never observed this in practice — for any sensible
            // lens distortion the Jacobian is dominated by the identity. If
            // it ever fires, the latest iterate is still the best estimate.
            debug_assert!(det.abs() > 1e-15, "singular Jacobian in undistort Newton step");
            if det.abs() < 1e-15 {
                break;
            }
            let inv_det = 1.0 / det;

            // Newton step: (x, y) ← (x, y) − J⁻¹ · r
            x -= inv_det * (j22 * rx - j12 * ry);
            y -= inv_det * (-j21 * rx + j11 * ry);
        }

        (x, y)
    }

    /// Returns `true` if all coefficients are zero.
    pub fn is_zero(&self) -> bool {
        self.a_coeffs.iter().all(|&c| c == 0.0)
            && self.b_coeffs.iter().all(|&c| c == 0.0)
            && self.ap_coeffs.iter().all(|&c| c == 0.0)
            && self.bp_coeffs.iter().all(|&c| c == 0.0)
    }
}

// ── Polynomial term helpers ─────────────────────────────────────────────────

/// Number of polynomial coefficients for the given order.
///
/// Terms are (p, q) with 0 ≤ p+q ≤ order:
///   order 0: 1 term   (0,0)  — constant offset (optical center shift)
///   order 1: +2 terms (1,0),(0,1)  — linear scale/rotation
///   order 2: +3 terms (2,0),(1,1),(0,2)  — quadratic
///   order 3: +4 terms (3,0),(2,1),(1,2),(0,3)  — cubic
///   etc.
///
/// Total = (order+1)(order+2)/2.
pub fn num_coeffs(order: u32) -> usize {
    ((order + 1) * (order + 2) / 2) as usize
}

/// Map (p, q) with 0 ≤ p+q ≤ order to a flat index.
///
/// Terms are enumerated in order of increasing sum, then decreasing p:
///   sum=0: (0,0)=0
///   sum=1: (1,0)=1, (0,1)=2
///   sum=2: (2,0)=3, (1,1)=4, (0,2)=5
///   sum=3: (3,0)=6, (2,1)=7, (1,2)=8, (0,3)=9
pub fn coeff_index(p: u32, q: u32) -> usize {
    let s = p + q;
    // Base offset: number of terms for sums 0..(s-1) = s*(s+1)/2
    let base = (s * (s + 1) / 2) as usize;
    // Within sum=s, terms are ordered by decreasing p: (s,0), (s-1,1), ..., (0,s)
    base + (s - p) as usize
}

/// Enumerate all (p, q) pairs for the given order.
pub fn term_pairs(order: u32) -> Vec<(u32, u32)> {
    let mut pairs = Vec::with_capacity(num_coeffs(order));
    for s in 0..=order {
        for p in (0..=s).rev() {
            let q = s - p;
            pairs.push((p, q));
        }
    }
    pairs
}

/// Evaluate a polynomial correction: Σ c_i · x^p_i · y^q_i
/// `coeffs` is a flat vector indexed by `coeff_index(p, q)`.
fn eval_poly(coeffs: &[f64], order: u32, x: f64, y: f64) -> f64 {
    let mut result = 0.0;
    let mut idx = 0;
    for s in 0..=order {
        for p in (0..=s).rev() {
            let q = s - p;
            result += coeffs[idx] * x.powi(p as i32) * y.powi(q as i32);
            idx += 1;
        }
    }
    result
}

/// Evaluate the polynomial `f(x, y) = Σ c_i · x^p · y^q` together with its
/// partial derivatives `(∂f/∂x, ∂f/∂y)`.
///
/// Returns `(value, df_dx, df_dy)`. Used by `PolynomialDistortion::undistort`
/// for Newton iteration on the forward polynomial.
fn eval_poly_with_grad(coeffs: &[f64], order: u32, x: f64, y: f64) -> (f64, f64, f64) {
    let mut value = 0.0;
    let mut df_dx = 0.0;
    let mut df_dy = 0.0;
    let mut idx = 0;
    for s in 0..=order {
        for p in (0..=s).rev() {
            let q = s - p;
            let c = coeffs[idx];
            let xp = x.powi(p as i32);
            let yq = y.powi(q as i32);
            value += c * xp * yq;
            // ∂(x^p · y^q)/∂x = p · x^(p-1) · y^q  (zero when p == 0)
            if p > 0 {
                df_dx += c * (p as f64) * x.powi(p as i32 - 1) * yq;
            }
            // ∂(x^p · y^q)/∂y = q · x^p · y^(q-1)  (zero when q == 0)
            if q > 0 {
                df_dy += c * (q as f64) * xp * y.powi(q as i32 - 1);
            }
            idx += 1;
        }
    }
    (value, df_dx, df_dy)
}

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

    #[test]
    fn test_num_coeffs() {
        assert_eq!(num_coeffs(0), 1); // (0,0)
        assert_eq!(num_coeffs(1), 3); // + (1,0),(0,1)
        assert_eq!(num_coeffs(2), 6); // + (2,0),(1,1),(0,2)
        assert_eq!(num_coeffs(3), 10); // + (3,0),(2,1),(1,2),(0,3)
        assert_eq!(num_coeffs(4), 15);
        assert_eq!(num_coeffs(5), 21);
    }

    #[test]
    fn test_coeff_index() {
        // sum=0
        assert_eq!(coeff_index(0, 0), 0);
        // sum=1
        assert_eq!(coeff_index(1, 0), 1);
        assert_eq!(coeff_index(0, 1), 2);
        // sum=2
        assert_eq!(coeff_index(2, 0), 3);
        assert_eq!(coeff_index(1, 1), 4);
        assert_eq!(coeff_index(0, 2), 5);
        // sum=3
        assert_eq!(coeff_index(3, 0), 6);
        assert_eq!(coeff_index(2, 1), 7);
        assert_eq!(coeff_index(1, 2), 8);
        assert_eq!(coeff_index(0, 3), 9);
        // sum=4
        assert_eq!(coeff_index(4, 0), 10);
        assert_eq!(coeff_index(3, 1), 11);
        assert_eq!(coeff_index(2, 2), 12);
        assert_eq!(coeff_index(1, 3), 13);
        assert_eq!(coeff_index(0, 4), 14);
    }

    #[test]
    fn test_term_pairs() {
        let pairs = term_pairs(3);
        assert_eq!(
            pairs,
            vec![
                (0, 0),
                (1, 0),
                (0, 1),
                (2, 0),
                (1, 1),
                (0, 2),
                (3, 0),
                (2, 1),
                (1, 2),
                (0, 3)
            ]
        );
    }

    #[test]
    fn test_zero_distortion_roundtrip() {
        let d = PolynomialDistortion::zero(4, 1024.0);
        let (xu, yu) = d.undistort(100.0, -200.0);
        assert!((xu - 100.0).abs() < 1e-12);
        assert!((yu + 200.0).abs() < 1e-12);
        let (xd, yd) = d.distort(100.0, -200.0);
        assert!((xd - 100.0).abs() < 1e-12);
        assert!((yd + 200.0).abs() < 1e-12);
    }

    #[test]
    fn test_newton_undistort_exact_inverse() {
        // Build a non-trivial forward distortion (no inverse coeffs set).
        // Newton iteration should still recover the ideal pixel exactly.
        let n = num_coeffs(4);
        let mut a = vec![0.0; n];
        let mut b = vec![0.0; n];
        a[coeff_index(2, 0)] = 0.01;
        a[coeff_index(0, 2)] = 0.005;
        a[coeff_index(1, 1)] = -0.003;
        b[coeff_index(2, 0)] = -0.004;
        b[coeff_index(0, 2)] = 0.012;
        b[coeff_index(3, 0)] = 0.001;

        let d = PolynomialDistortion::new(4, 1024.0, a, b, vec![0.0; n], vec![0.0; n]);

        // Forward then back must roundtrip to within numerical precision.
        for &(x, y) in &[(0.0, 0.0), (100.0, -200.0), (500.0, 400.0), (-800.0, 100.0)] {
            let (xd, yd) = d.distort(x, y);
            let (xu, yu) = d.undistort(xd, yd);
            assert!(
                (xu - x).abs() < 1e-9,
                "x roundtrip {} -> {} -> {} (err {:.2e})",
                x,
                xd,
                xu,
                xu - x
            );
            assert!(
                (yu - y).abs() < 1e-9,
                "y roundtrip {} -> {} -> {} (err {:.2e})",
                y,
                yd,
                yu,
                yu - y
            );
        }
    }

    #[test]
    fn test_distort_undistort_basic() {
        // Create a simple distortion: only x² and y² terms
        let n = num_coeffs(4);
        let mut a = vec![0.0; n];
        let mut b = vec![0.0; n];
        a[coeff_index(2, 0)] = 0.01; // x² → dx
        b[coeff_index(0, 2)] = -0.005; // y² → dy

        // Compute inverse from a grid (simple test)
        let d = PolynomialDistortion::new(
            4,
            1024.0,
            a,
            b,
            vec![0.0; n], // ap (inverse) not set here
            vec![0.0; n], // bp
        );

        // Forward: distort(0, 0) = (0, 0)
        let (xd, yd) = d.distort(0.0, 0.0);
        assert!(xd.abs() < 1e-12);
        assert!(yd.abs() < 1e-12);

        // Forward at (512, 512): u=0.5, v=0.5
        // dx = 0.01 * 0.5² = 0.0025 (normalized), * 1024 = 2.56 px
        let (xd, _yd) = d.distort(512.0, 512.0);
        assert!((xd - 512.0 - 2.56).abs() < 1e-10);
    }
}