tx2_iff/

warp.rs

//! Symplectic warping field (Layer 3: The Warping Field)
//!
//! This module implements symplectic spatial advection using 4-fold radial
//! symmetry and vortex primitives. Instead of storing texture for every fold
//! in fabric, we store the texture of a "flat" patch and a vector field that
//! describes how to warp it.
//!
//! ## Mathematical Foundation
//!
//! The warping field preserves area (incompressibility):
//! ```text
//! ∇·u = 0  (divergence-free constraint)
//! ```
//!
//! Implemented via stream function:
//! ```text
//! u = ∂ψ/∂y
//! v = -∂ψ/∂x
//! ```
//!
//! ## 4-Fold Radial Basis
//!
//! Flow field decomposition:
//! ```text
//! u(x,y) = ∑ₖ₌₁⁴ cₖ ψₖ(r,θ)
//! ```
//!
//! Where:
//! - ψ₁(r,θ) = r cos(θ)     - Radial outward
//! - ψ₂(r,θ) = r sin(θ)     - Radial rotated 90°
//! - ψ₃(r,θ) = r cos(4θ)    - 4-fold symmetric
//! - ψ₄(r,θ) = r sin(4θ)    - 4-fold antisymmetric

use crate::fixed::Fixed;
use serde::{Deserialize, Serialize};

/// Vortex primitive for local warping
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Vortex {
    /// Center position X
    pub x: u16,
    /// Center position Y
    pub y: u16,
    /// Angular velocity (signed for rotation direction)
    pub strength: i16,
    /// Influence radius
    pub radius: u8,
    /// Exponential decay rate
    pub decay: u8,
}

impl Vortex {
    /// Create a new vortex
    pub fn new(x: u16, y: u16, strength: i16, radius: u8, decay: u8) -> Self {
        Vortex {
            x,
            y,
            strength,
            radius,
            decay,
        }
    }

    /// Calculate velocity field contribution at point (x, y)
    pub fn velocity_at(&self, x: Fixed, y: Fixed) -> (Fixed, Fixed) {
        // Vector from vortex center to point
        let dx = x - Fixed::from(self.x);
        let dy = y - Fixed::from(self.y);

        // Distance from center
        let r_sq = dx * dx + dy * dy;
        let r = r_sq.sqrt().unwrap_or(Fixed::ZERO);

        // Avoid singularity at center
        if r < Fixed::ONE {
            return (Fixed::ZERO, Fixed::ZERO);
        }

        // Angle
        let theta = atan2_fixed(dy, dx);

        // Falloff function: exp(-decay * r / radius)
        let decay_fixed = Fixed::from_f32(self.decay as f32 / 255.0);
        let radius_fixed = Fixed::from_int(self.radius as i32);
        let falloff_arg = -(decay_fixed * r / radius_fixed);
        let falloff = exp_fixed(falloff_arg);

        // Vortex strength
        let strength_fixed = Fixed::from_f32(self.strength as f32 / 100.0);

        // Velocity field for vortex (perpendicular to radial direction)
        // u = strength × sin(θ) × falloff
        // v = -strength × cos(θ) × falloff
        let u = strength_fixed * theta.sin() * falloff;
        let v = -strength_fixed * theta.cos() * falloff;

        (u, v)
    }
}

/// 4-fold radial basis coefficients
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct RadialBasis {
    /// Coefficients for the 4 basis functions
    pub coeffs: [Fixed; 4],
}

impl RadialBasis {
    /// Create new radial basis with zero coefficients
    pub fn zero() -> Self {
        RadialBasis {
            coeffs: [Fixed::ZERO; 4],
        }
    }

    /// Create from raw coefficients
    pub fn new(c1: Fixed, c2: Fixed, c3: Fixed, c4: Fixed) -> Self {
        RadialBasis {
            coeffs: [c1, c2, c3, c4],
        }
    }

    /// Evaluate flow field at point (x, y) relative to center
    pub fn velocity_at(&self, x: Fixed, y: Fixed, center_x: Fixed, center_y: Fixed) -> (Fixed, Fixed) {
        // Relative coordinates
        let dx = x - center_x;
        let dy = y - center_y;

        // Polar coordinates
        let r_sq = dx * dx + dy * dy;
        let r = r_sq.sqrt().unwrap_or(Fixed::ZERO);
        let theta = atan2_fixed(dy, dx);

        // Four-fold symmetric angle
        let theta_4 = theta * Fixed::from_int(4);

        // Basis functions
        // ψ₁(r,θ) = r cos(θ)
        let psi1_x = r * theta.cos();
        let psi1_y = Fixed::ZERO;

        // ψ₂(r,θ) = r sin(θ)
        let psi2_x = Fixed::ZERO;
        let psi2_y = r * theta.sin();

        // ψ₃(r,θ) = r cos(4θ)
        let psi3_x = r * theta_4.cos();
        let psi3_y = Fixed::ZERO;

        // ψ₄(r,θ) = r sin(4θ)
        let psi4_x = Fixed::ZERO;
        let psi4_y = r * theta_4.sin();

        // Combine with coefficients
        let u = self.coeffs[0] * psi1_x + self.coeffs[1] * psi2_x
            + self.coeffs[2] * psi3_x + self.coeffs[3] * psi4_x;
        let v = self.coeffs[0] * psi1_y + self.coeffs[1] * psi2_y
            + self.coeffs[2] * psi3_y + self.coeffs[3] * psi4_y;

        (u, v)
    }
}

/// Warp field combining global radial basis and local vortices
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WarpField {
    /// Image dimensions
    pub width: u32,
    pub height: u32,
    /// Global radial basis (4 coefficients)
    pub global_basis: RadialBasis,
    /// Center point for global basis
    pub center_x: u16,
    pub center_y: u16,
    /// Local vortex primitives
    pub vortices: Vec<Vortex>,
}

impl WarpField {
    /// Create new warp field
    pub fn new(width: u32, height: u32) -> Self {
        WarpField {
            width,
            height,
            global_basis: RadialBasis::zero(),
            center_x: (width / 2) as u16,
            center_y: (height / 2) as u16,
            vortices: Vec::new(),
        }
    }

    /// Add a vortex
    pub fn add_vortex(&mut self, vortex: Vortex) {
        self.vortices.push(vortex);
    }

    /// Calculate total velocity field at point (x, y)
    pub fn velocity_at(&self, x: u16, y: u16) -> (Fixed, Fixed) {
        let x_fixed = Fixed::from(x);
        let y_fixed = Fixed::from(y);

        // Global contribution
        let (u_global, v_global) = self.global_basis.velocity_at(
            x_fixed,
            y_fixed,
            Fixed::from(self.center_x),
            Fixed::from(self.center_y),
        );

        // Local vortex contributions
        let mut u_total = u_global;
        let mut v_total = v_global;

        for vortex in &self.vortices {
            let (u_vortex, v_vortex) = vortex.velocity_at(x_fixed, y_fixed);
            u_total = u_total + u_vortex;
            v_total = v_total + v_vortex;
        }

        (u_total, v_total)
    }

    /// Apply warping to get source coordinates (backtracing)
    pub fn warp_backwards(&self, x: u16, y: u16) -> (Fixed, Fixed) {
        let (u, v) = self.velocity_at(x, y);

        // Source position = current position - velocity
        let x_fixed = Fixed::from(x);
        let y_fixed = Fixed::from(y);

        (x_fixed - u, y_fixed - v)
    }
}

/// Bicubic interpolation for warped sampling
pub struct BicubicSampler;

impl BicubicSampler {
    /// Sample image at fixed-point coordinates using bicubic interpolation
    pub fn sample(image: &[[u8; 3]], width: usize, height: usize, x: Fixed, y: Fixed) -> [u8; 3] {
        // Get integer coordinates
        let xi = x.floor().to_int().clamp(0, width as i32 - 1) as usize;
        let yi = y.floor().to_int().clamp(0, height as i32 - 1) as usize;

        // Get fractional parts
        let xf = x.fract();
        let yf = y.fract();

        // Sample 4x4 neighborhood
        let mut pixels = [[0u8; 3]; 16];
        for dy in 0..4 {
            for dx in 0..4 {
                let sample_x = (xi + dx).saturating_sub(1).min(width - 1);
                let sample_y = (yi + dy).saturating_sub(1).min(height - 1);
                let idx = sample_y * width + sample_x;
                if idx < image.len() {
                    pixels[dy * 4 + dx] = image[idx];
                }
            }
        }

        // Bicubic weights: [-1, 9, 9, -1] / 16 for each dimension
        Self::bicubic_interpolate(&pixels, xf, yf)
    }

    /// Bicubic interpolation kernel
    fn bicubic_interpolate(pixels: &[[u8; 3]; 16], xf: Fixed, yf: Fixed) -> [u8; 3] {
        let mut result = [0u8; 3];

        // Weights for bicubic kernel
        let weights = Self::cubic_weights(xf);
        let weights_y = Self::cubic_weights(yf);

        for c in 0..3 {
            let mut sum = Fixed::ZERO;

            for y in 0..4 {
                let mut row_sum = Fixed::ZERO;
                for x in 0..4 {
                    let pixel_val = Fixed::from_f32(pixels[y * 4 + x][c] as f32);
                    row_sum = row_sum + pixel_val * weights[x];
                }
                sum = sum + row_sum * weights_y[y];
            }

            result[c] = sum.to_int().clamp(0, 255) as u8;
        }

        result
    }

    /// Calculate cubic interpolation weights
    fn cubic_weights(t: Fixed) -> [Fixed; 4] {
        // Catmull-Rom spline weights
        let t2 = t * t;
        let t3 = t2 * t;

        let half = Fixed::HALF;

        [
            -half * t3 + t2 - half * t,
            Fixed::from_f32(1.5) * t3 - Fixed::from_f32(2.5) * t2 + Fixed::ONE,
            -Fixed::from_f32(1.5) * t3 + Fixed::from_int(2) * t2 + half * t,
            half * t3 - half * t2,
        ]
    }
}

/// Fixed-point atan2 approximation
fn atan2_fixed(y: Fixed, x: Fixed) -> Fixed {
    // Special cases
    if x == Fixed::ZERO {
        if y > Fixed::ZERO {
            return Fixed::PI / Fixed::from_int(2);
        } else if y < Fixed::ZERO {
            return -Fixed::PI / Fixed::from_int(2);
        } else {
            return Fixed::ZERO;
        }
    }

    // Use atan(y/x) approximation
    let ratio = y / x;
    let atan_val = atan_approx(ratio);

    // Adjust for quadrant
    if x > Fixed::ZERO {
        atan_val
    } else if y >= Fixed::ZERO {
        atan_val + Fixed::PI
    } else {
        atan_val - Fixed::PI
    }
}

/// Fixed-point atan approximation using polynomial
fn atan_approx(x: Fixed) -> Fixed {
    // atan(x) ≈ x - x³/3 + x⁵/5 - x⁷/7 (for small x)
    // For larger x, use atan(x) = π/2 - atan(1/x)

    let abs_x = x.abs();

    if abs_x > Fixed::ONE {
        // Use identity: atan(x) = π/2 - atan(1/x)
        let recip = Fixed::ONE / abs_x;
        let atan_recip = atan_approx(recip);
        let result = Fixed::PI / Fixed::from_int(2) - atan_recip;

        if x < Fixed::ZERO {
            -result
        } else {
            result
        }
    } else {
        // Taylor series
        let x2 = x * x;
        let x3 = x2 * x;
        let x5 = x3 * x2;
        let x7 = x5 * x2;

        x - x3 / Fixed::from_int(3) + x5 / Fixed::from_int(5) - x7 / Fixed::from_int(7)
    }
}

/// Fixed-point exponential approximation
fn exp_fixed(x: Fixed) -> Fixed {
    // Clamp to reasonable range
    let x_clamped = x.clamp(Fixed::from_int(-10), Fixed::from_int(10));

    // For negative values, use exp(-x) = 1/exp(x)
    if x_clamped < Fixed::ZERO {
        let pos_exp = exp_fixed_positive(-x_clamped);
        // Avoid division by zero
        if pos_exp == Fixed::ZERO {
            return Fixed::ZERO;
        }
        return Fixed::ONE / pos_exp;
    }

    exp_fixed_positive(x_clamped)
}

/// Helper for positive exponent using Taylor series
fn exp_fixed_positive(x: Fixed) -> Fixed {
    // exp(x) ≈ 1 + x + x²/2 + x³/6 + x⁴/24 + x⁵/120
    let x2 = x * x;
    let x3 = x2 * x;
    let x4 = x3 * x;
    let x5 = x4 * x;

    Fixed::ONE + x
        + x2 / Fixed::from_int(2)
        + x3 / Fixed::from_int(6)
        + x4 / Fixed::from_int(24)
        + x5 / Fixed::from_int(120)
}

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

    #[test]
    fn test_vortex_velocity() {
        let vortex = Vortex::new(100, 100, 100, 50, 128);

        // At center, velocity should be near zero
        let (u, v) = vortex.velocity_at(Fixed::from(100), Fixed::from(100));
        assert!(u.abs() < Fixed::from_int(1));
        assert!(v.abs() < Fixed::from_int(1));

        // Away from center, should have non-zero velocity
        let (u, v) = vortex.velocity_at(Fixed::from(120), Fixed::from(100));
        assert!(u.abs() > Fixed::ZERO || v.abs() > Fixed::ZERO);
    }

    #[test]
    fn test_radial_basis() {
        let basis = RadialBasis::new(
            Fixed::from_int(1),
            Fixed::ZERO,
            Fixed::ZERO,
            Fixed::ZERO,
        );

        let center = Fixed::from(50);
        let (u, v) = basis.velocity_at(
            Fixed::from(60),
            Fixed::from(50),
            center,
            center,
        );

        // With only first coefficient, should have radial outward flow
        assert!(u > Fixed::ZERO);
    }

    #[test]
    fn test_warp_field() {
        let mut field = WarpField::new(200, 200);

        field.add_vortex(Vortex::new(100, 100, 50, 30, 128));

        let (u, v) = field.velocity_at(120, 100);

        // Should have velocity from vortex
        assert!(u.abs() > Fixed::ZERO || v.abs() > Fixed::ZERO);
    }

    #[test]
    fn test_warp_backwards() {
        let mut field = WarpField::new(200, 200);

        // Add a vortex with proper parameters (lower decay = wider influence)
        field.add_vortex(Vortex::new(100, 100, 500, 50, 50));

        let (source_x, source_y) = field.warp_backwards(120, 100);

        // Source should be different from target due to vortex warping
        let target_x = Fixed::from(120);
        let target_y = Fixed::from(100);

        let dx = if source_x > target_x {
            source_x - target_x
        } else {
            target_x - source_x
        };

        let dy = if source_y > target_y {
            source_y - target_y
        } else {
            target_y - source_y
        };

        // At least one dimension should show warp effect
        assert!(
            dx > Fixed::from_f32(0.1) || dy > Fixed::from_f32(0.1),
            "Expected warp to have visible effect, dx={}, dy={}",
            dx.to_f32(),
            dy.to_f32()
        );
    }

    #[test]
    fn test_bicubic_sampler() {
        // Create a simple test image
        let width = 10;
        let height = 10;
        let mut image = vec![[100u8; 3]; width * height];

        // Set a few pixels to different colors
        image[0] = [255, 0, 0];
        image[width - 1] = [0, 255, 0];
        image[(height - 1) * width] = [0, 0, 255];

        // Sample at integer coordinate
        let pixel = BicubicSampler::sample(
            &image,
            width,
            height,
            Fixed::from_int(0),
            Fixed::from_int(0),
        );
        assert_eq!(pixel[0], 255);

        // Sample at fractional coordinate
        let pixel = BicubicSampler::sample(
            &image,
            width,
            height,
            Fixed::from_f32(0.5),
            Fixed::from_f32(0.5),
        );
        // Should be interpolated value
        assert!(pixel[0] > 100 && pixel[0] < 255);
    }

    #[test]
    fn test_atan2_approximation() {
        let test_cases = vec![
            (1.0, 1.0, std::f32::consts::PI / 4.0),      // atan2(1, 1) = π/4
            (1.0, 0.0, std::f32::consts::PI / 2.0),      // atan2(1, 0) = π/2
            (0.0, 1.0, 0.0),                              // atan2(0, 1) = 0
            (-1.0, 1.0, -std::f32::consts::PI / 4.0),    // atan2(-1, 1) = -π/4
        ];

        for (y, x, expected) in test_cases {
            let y_fixed = Fixed::from_f32(y);
            let x_fixed = Fixed::from_f32(x);
            let result = atan2_fixed(y_fixed, x_fixed);

            // Allow 10% error in approximation
            assert!(
                (result.to_f32() - expected).abs() < 0.1,
                "atan2({}, {}) = {}, expected {}",
                y,
                x,
                result.to_f32(),
                expected
            );
        }
    }

    #[test]
    fn test_exp_approximation() {
        let test_cases = vec![
            (0.0, 1.0),
            (1.0, std::f32::consts::E),
            (-1.0, 1.0 / std::f32::consts::E),
            (2.0, std::f32::consts::E * std::f32::consts::E),
        ];

        for (input, expected) in test_cases {
            let input_fixed = Fixed::from_f32(input);
            let result = exp_fixed(input_fixed);

            // Allow 10% error
            let error = (result.to_f32() - expected).abs() / expected;
            assert!(error < 0.1, "exp({}) = {}, expected {}", input, result.to_f32(), expected);
        }
    }

    #[test]
    fn test_four_fold_symmetry() {
        let basis = RadialBasis::new(
            Fixed::ZERO,
            Fixed::ZERO,
            Fixed::ONE,
            Fixed::ZERO,
        );

        let center = Fixed::from(50);

        // Test at 4 symmetric points (every 90 degrees)
        let points = [
            (60, 50), // 0 degrees
            (50, 60), // 90 degrees
            (40, 50), // 180 degrees
            (50, 40), // 270 degrees
        ];

        let mut velocities = Vec::new();
        for (x, y) in &points {
            let (u, v) = basis.velocity_at(
                Fixed::from(*x),
                Fixed::from(*y),
                center,
                center,
            );
            velocities.push((u, v));
        }

        // With 4-fold symmetry, magnitudes should be similar
        let magnitudes: Vec<Fixed> = velocities
            .iter()
            .map(|(u, v)| (*u * *u + *v * *v).sqrt().unwrap_or(Fixed::ZERO))
            .collect();

        for i in 1..magnitudes.len() {
            let ratio = magnitudes[i] / magnitudes[0];
            // Should be within 20% of each other
            assert!(ratio > Fixed::from_f32(0.8) && ratio < Fixed::from_f32(1.2));
        }
    }
}