//! Geometry encoding for neural operators.

use aeroflow_shared::{AirfoilGeometry, Point2D, WingGeometry};

/// Geometry encoder for converting airfoil/wing shapes to feature vectors.
#[derive(Debug)]
pub struct GeometryEncoder {
    /// Hidden dimension for encoded features.
    hidden_dim: usize,
    /// Number of Chebyshev nodes for interpolation.
    #[allow(dead_code)]
    num_nodes: usize,
}

impl GeometryEncoder {
    /// Create a new geometry encoder.
    pub fn new(hidden_dim: usize) -> Self {
        Self {
            hidden_dim,
            num_nodes: 64,
        }
    }

    /// Encode an airfoil geometry to feature vector.
    pub fn encode_airfoil(&self, airfoil: &AirfoilGeometry) -> Vec<f32> {
        let mut features = Vec::with_capacity(self.hidden_dim);

        // Global shape parameters
        features.push(airfoil.max_thickness);
        features.push(airfoil.max_camber);
        features.push(airfoil.max_camber_position);
        features.push(airfoil.chord);

        // Encode surface shape using Chebyshev coefficients
        let upper_coeffs = self.chebyshev_coefficients(&airfoil.upper_surface, 10);
        let lower_coeffs = self.chebyshev_coefficients(&airfoil.lower_surface, 10);

        features.extend(upper_coeffs);
        features.extend(lower_coeffs);

        // Curvature features
        let upper_curvature = self.compute_curvature(&airfoil.upper_surface);
        let lower_curvature = self.compute_curvature(&airfoil.lower_surface);
        features.push(upper_curvature.iter().sum::<f32>() / upper_curvature.len() as f32);
        features.push(lower_curvature.iter().sum::<f32>() / lower_curvature.len() as f32);

        // Leading edge radius
        features.push(self.estimate_le_radius(airfoil));

        // Trailing edge angle
        features.push(self.compute_te_angle(airfoil));

        // Pad or truncate to hidden_dim
        features.resize(self.hidden_dim, 0.0);
        features
    }

    /// Encode a wing geometry to feature vector.
    pub fn encode_wing(&self, wing: &WingGeometry) -> Vec<f32> {
        let mut features = Vec::with_capacity(self.hidden_dim);

        // Encode root airfoil
        let root_features = self.encode_airfoil(&wing.root_airfoil);
        features.extend(root_features.iter().take(20));

        // Planform parameters
        features.push(wing.span);
        features.push(wing.root_chord);
        features.push(wing.tip_chord);
        features.push(wing.tip_chord / wing.root_chord); // Taper ratio
        features.push(wing.sweep_angle.to_radians());
        features.push(wing.dihedral_angle.to_radians());
        features.push(wing.twist_angle.to_radians());

        // Aspect ratio
        let avg_chord = f32::midpoint(wing.root_chord, wing.tip_chord);
        let wing_area = wing.span * avg_chord;
        let aspect_ratio = wing.span.powi(2) / wing_area;
        features.push(aspect_ratio);

        // Pad to hidden_dim
        features.resize(self.hidden_dim, 0.0);
        features
    }

    /// Compute Chebyshev coefficients for surface representation.
    fn chebyshev_coefficients(&self, points: &[Point2D], num_coeffs: usize) -> Vec<f32> {
        if points.is_empty() {
            return vec![0.0; num_coeffs];
        }

        let n = points.len();
        let mut coeffs = vec![0.0; num_coeffs];

        for k in 0..num_coeffs {
            let mut sum = 0.0;
            for (j, point) in points.iter().enumerate() {
                let x = 2.0 * j as f32 / (n - 1) as f32 - 1.0; // Map to [-1, 1]
                let tk = self.chebyshev_t(k, x);
                sum += point.y * tk;
            }
            coeffs[k] = sum / n as f32;
        }

        coeffs
    }

    /// Chebyshev polynomial of the first kind.
    fn chebyshev_t(&self, n: usize, x: f32) -> f32 {
        match n {
            0 => 1.0,
            1 => x,
            _ => 2.0 * x * self.chebyshev_t(n - 1, x) - self.chebyshev_t(n - 2, x),
        }
    }

    /// Compute curvature along a surface.
    fn compute_curvature(&self, points: &[Point2D]) -> Vec<f32> {
        if points.len() < 3 {
            return vec![0.0; points.len()];
        }

        let mut curvature = vec![0.0; points.len()];

        for i in 1..points.len() - 1 {
            let p0 = &points[i - 1];
            let p1 = &points[i];
            let p2 = &points[i + 1];

            // First derivatives
            let dx1 = p1.x - p0.x;
            let dy1 = p1.y - p0.y;
            let dx2 = p2.x - p1.x;
            let dy2 = p2.y - p1.y;

            // Second derivatives
            let ddx = dx2 - dx1;
            let ddy = dy2 - dy1;

            // Curvature formula: κ = |x'y'' - y'x''| / (x'^2 + y'^2)^(3/2)
            let dx_avg = f32::midpoint(dx1, dx2);
            let dy_avg = f32::midpoint(dy1, dy2);
            let denom = (dx_avg.powi(2) + dy_avg.powi(2)).powf(1.5).max(1e-10);
            curvature[i] = (dx_avg * ddy - dy_avg * ddx).abs() / denom;
        }

        curvature
    }

    /// Estimate leading edge radius.
    fn estimate_le_radius(&self, airfoil: &AirfoilGeometry) -> f32 {
        // Approximate LE radius from thickness distribution
        // For NACA 4-digit: r_LE ≈ 1.1019 * (t/c)^2
        1.1019 * airfoil.max_thickness.powi(2)
    }

    /// Compute trailing edge angle.
    fn compute_te_angle(&self, airfoil: &AirfoilGeometry) -> f32 {
        if airfoil.upper_surface.len() < 2 || airfoil.lower_surface.len() < 2 {
            return 0.0;
        }

        let upper_last = &airfoil.upper_surface[airfoil.upper_surface.len() - 1];
        let upper_prev = &airfoil.upper_surface[airfoil.upper_surface.len() - 2];
        let lower_last = &airfoil.lower_surface[airfoil.lower_surface.len() - 1];
        let lower_prev = &airfoil.lower_surface[airfoil.lower_surface.len() - 2];

        let upper_slope = (upper_last.y - upper_prev.y) / (upper_last.x - upper_prev.x).max(1e-6);
        let lower_slope = (lower_last.y - lower_prev.y) / (lower_last.x - lower_prev.x).max(1e-6);

        // TE angle is the angle between upper and lower surface slopes
        (upper_slope.atan() - lower_slope.atan()).abs()
    }

    /// Interpolate surface at given x positions.
    pub fn interpolate_surface(&self, points: &[Point2D], x_positions: &[f32]) -> Vec<f32> {
        x_positions
            .iter()
            .map(|&x| {
                // Find bracketing points
                let mut y = 0.0;
                for i in 1..points.len() {
                    if points[i].x >= x {
                        let t = (x - points[i - 1].x) / (points[i].x - points[i - 1].x).max(1e-10);
                        y = points[i - 1].y + t * (points[i].y - points[i - 1].y);
                        break;
                    }
                }
                y
            })
            .collect()
    }

    /// Compute area enclosed by airfoil.
    pub fn compute_area(&self, airfoil: &AirfoilGeometry) -> f32 {
        let mut area = 0.0;

        // Shoelace formula for upper surface
        for i in 1..airfoil.upper_surface.len() {
            let p0 = &airfoil.upper_surface[i - 1];
            let p1 = &airfoil.upper_surface[i];
            area += (p1.x - p0.x) * (p0.y + p1.y) / 2.0;
        }

        // Subtract lower surface (traversed in reverse)
        for i in 1..airfoil.lower_surface.len() {
            let p0 = &airfoil.lower_surface[i - 1];
            let p1 = &airfoil.lower_surface[i];
            area -= (p1.x - p0.x) * (p0.y + p1.y) / 2.0;
        }

        area.abs()
    }
}

/// CST (Class-Shape Transformation) parameterization.
#[derive(Debug, Clone)]
pub struct CSTAirfoil {
    /// Class function exponents.
    pub n1: f32,
    pub n2: f32,
    /// Upper surface coefficients.
    pub upper_coeffs: Vec<f32>,
    /// Lower surface coefficients.
    pub lower_coeffs: Vec<f32>,
    /// Trailing edge thickness.
    pub te_thickness: f32,
}

impl Default for CSTAirfoil {
    fn default() -> Self {
        Self {
            n1: 0.5,
            n2: 1.0,
            upper_coeffs: vec![0.2, 0.3, 0.2, 0.1, 0.1],
            lower_coeffs: vec![-0.2, -0.1, -0.1, -0.05, -0.05],
            te_thickness: 0.0,
        }
    }
}

impl CSTAirfoil {
    /// Generate airfoil geometry from CST parameters.
    pub fn to_airfoil(&self, num_points: usize) -> AirfoilGeometry {
        let mut upper = Vec::with_capacity(num_points);
        let mut lower = Vec::with_capacity(num_points);

        for i in 0..num_points {
            let psi = 1.0 - (std::f32::consts::PI * i as f32 / (num_points - 1) as f32).cos();
            let psi = psi / 2.0; // Map to [0, 1]

            // Class function
            let class_func = psi.powf(self.n1) * (1.0 - psi).powf(self.n2);

            // Shape functions (Bernstein polynomials)
            let upper_shape = self.bernstein_sum(&self.upper_coeffs, psi);
            let lower_shape = self.bernstein_sum(&self.lower_coeffs, psi);

            let y_upper = class_func * upper_shape + psi * self.te_thickness / 2.0;
            let y_lower = class_func * lower_shape - psi * self.te_thickness / 2.0;

            upper.push(Point2D::new(psi, y_upper));
            lower.push(Point2D::new(psi, y_lower));
        }

        let max_thickness = upper
            .iter()
            .zip(lower.iter())
            .map(|(u, l)| u.y - l.y)
            .fold(0.0f32, f32::max);

        AirfoilGeometry {
            name: "CST Airfoil".to_string(),
            upper_surface: upper,
            lower_surface: lower,
            chord: 1.0,
            max_thickness,
            max_camber: 0.0,
            max_camber_position: 0.0,
        }
    }

    /// Evaluate Bernstein polynomial sum.
    fn bernstein_sum(&self, coeffs: &[f32], t: f32) -> f32 {
        let n = coeffs.len() - 1;
        let mut sum = 0.0;

        for (k, &coeff) in coeffs.iter().enumerate() {
            sum += coeff * self.bernstein(n, k, t);
        }

        sum
    }

    /// Bernstein basis polynomial.
    fn bernstein(&self, n: usize, k: usize, t: f32) -> f32 {
        let binom = self.binomial(n, k) as f32;
        binom * t.powi(k as i32) * (1.0 - t).powi((n - k) as i32)
    }

    /// Binomial coefficient.
    fn binomial(&self, n: usize, k: usize) -> usize {
        if k > n {
            return 0;
        }
        let mut result = 1;
        for i in 0..k {
            result = result * (n - i) / (i + 1);
        }
        result
    }
}

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

    #[test]
    fn test_encoder_creation() {
        let encoder = GeometryEncoder::new(64);
        assert_eq!(encoder.hidden_dim, 64);
    }

    #[test]
    fn test_encode_airfoil() {
        let encoder = GeometryEncoder::new(32);
        let airfoil = naca_4digit("0012", 50).unwrap();
        let features = encoder.encode_airfoil(&airfoil);

        assert_eq!(features.len(), 32);
        assert!((features[0] - 0.12).abs() < 0.001); // max_thickness
    }

    #[test]
    fn test_encode_wing() {
        let encoder = GeometryEncoder::new(64);
        let wing = WingGeometry::default();
        let features = encoder.encode_wing(&wing);

        assert_eq!(features.len(), 64);
    }

    #[test]
    fn test_chebyshev() {
        let encoder = GeometryEncoder::new(32);

        // T_0(x) = 1
        assert!((encoder.chebyshev_t(0, 0.5) - 1.0).abs() < 0.001);

        // T_1(x) = x
        assert!((encoder.chebyshev_t(1, 0.5) - 0.5).abs() < 0.001);

        // T_2(x) = 2x^2 - 1
        assert!((encoder.chebyshev_t(2, 0.5) - (-0.5)).abs() < 0.001);
    }

    #[test]
    fn test_curvature() {
        let encoder = GeometryEncoder::new(32);
        let points = vec![
            Point2D::new(0.0, 0.0),
            Point2D::new(0.5, 0.1),
            Point2D::new(1.0, 0.0),
        ];
        let curvature = encoder.compute_curvature(&points);
        assert_eq!(curvature.len(), 3);
        assert!(curvature[1] > 0.0); // Should have some curvature
    }

    #[test]
    fn test_area() {
        let encoder = GeometryEncoder::new(32);
        let airfoil = naca_4digit("0012", 50).unwrap();
        let area = encoder.compute_area(&airfoil);
        assert!(area > 0.0);
        assert!(area < 0.2); // Reasonable area for NACA 0012
    }

    #[test]
    fn test_cst_airfoil() {
        let cst = CSTAirfoil::default();
        let airfoil = cst.to_airfoil(50);

        assert_eq!(airfoil.upper_surface.len(), 50);
        assert_eq!(airfoil.lower_surface.len(), 50);
        assert!(airfoil.max_thickness > 0.0);
    }

    #[test]
    fn test_bernstein() {
        let cst = CSTAirfoil::default();

        // B_{0,0}(0.5) = 1
        assert!((cst.bernstein(0, 0, 0.5) - 1.0).abs() < 0.001);

        // B_{2,1}(0.5) = 2 * 0.5 * 0.5 = 0.5
        assert!((cst.bernstein(2, 1, 0.5) - 0.5).abs() < 0.001);
    }
}