//! Physics-Informed Neural Network for Structural Mechanics
//!
//! This module implements the neural network architecture and physics losses
//! for solving structural mechanics problems using PINNs.

use structural_shared::{BoundaryCondition, Bounds2D, DisplacementField, Material, StressField};

// ============================================================================
// Activation Functions
// ============================================================================

/// Activation function type.
#[derive(Debug, Clone, Copy)]
pub enum Activation {
    Tanh,
    Relu,
    Sigmoid,
    Swish,
    Gelu,
    Sin, // For Fourier-like feature embedding
}

impl Activation {
    /// Parse activation from string.
    pub fn from_str(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "tanh" => Self::Tanh,
            "relu" => Self::Relu,
            "sigmoid" => Self::Sigmoid,
            "swish" => Self::Swish,
            "gelu" => Self::Gelu,
            "sin" => Self::Sin,
            _ => Self::Tanh,
        }
    }

    /// Apply activation function.
    pub fn apply(&self, x: f64) -> f64 {
        match self {
            Self::Tanh => x.tanh(),
            Self::Relu => x.max(0.0),
            Self::Sigmoid => 1.0 / (1.0 + (-x).exp()),
            Self::Swish => x * (1.0 / (1.0 + (-x).exp())),
            Self::Gelu => {
                0.5 * x
                    * (1.0
                        + (std::f64::consts::FRAC_2_SQRT_PI * (x + 0.044715 * x.powi(3))
                            / std::f64::consts::SQRT_2)
                            .tanh())
            }
            Self::Sin => x.sin(),
        }
    }

    /// Apply derivative of activation function.
    pub fn derivative(&self, x: f64) -> f64 {
        match self {
            Self::Tanh => 1.0 - x.tanh().powi(2),
            Self::Relu => {
                if x > 0.0 {
                    1.0
                } else {
                    0.0
                }
            }
            Self::Sigmoid => {
                let s = 1.0 / (1.0 + (-x).exp());
                s * (1.0 - s)
            }
            Self::Swish => {
                let s = 1.0 / (1.0 + (-x).exp());
                s + x * s * (1.0 - s)
            }
            Self::Gelu => {
                // Approximate derivative
                let cdf = 0.5 * (1.0 + (x / std::f64::consts::SQRT_2).tanh());
                let pdf = (-0.5 * x * x).exp() / (2.0 * std::f64::consts::PI).sqrt();
                cdf + x * pdf
            }
            Self::Sin => x.cos(),
        }
    }
}

// ============================================================================
// Dense Layer
// ============================================================================

/// Dense (fully connected) layer.
#[derive(Debug, Clone)]
pub struct DenseLayer {
    /// Weight matrix (out_dim x in_dim).
    weights: Vec<Vec<f64>>,
    /// Bias vector.
    biases: Vec<f64>,
    /// Input dimension.
    in_dim: usize,
    /// Output dimension.
    out_dim: usize,
}

impl DenseLayer {
    /// Create a new dense layer with Xavier initialization.
    pub fn new(in_dim: usize, out_dim: usize, seed: u64) -> Self {
        let mut rng_state = seed;
        let scale = (2.0 / (in_dim + out_dim) as f64).sqrt();

        let mut weights = vec![vec![0.0; in_dim]; out_dim];
        let mut biases = vec![0.0; out_dim];

        // Xavier initialization
        for row in &mut weights {
            for w in row {
                rng_state = rng_state
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                let u = (rng_state >> 11) as f64 / (1u64 << 53) as f64;
                *w = (u * 2.0 - 1.0) * scale;
            }
        }

        for b in &mut biases {
            *b = 0.0; // Initialize biases to zero
        }

        Self {
            weights,
            biases,
            in_dim,
            out_dim,
        }
    }

    /// Forward pass.
    pub fn forward(&self, input: &[f64]) -> Vec<f64> {
        assert_eq!(input.len(), self.in_dim);

        let mut output = vec![0.0; self.out_dim];
        for i in 0..self.out_dim {
            let mut sum = self.biases[i];
            for j in 0..self.in_dim {
                sum += self.weights[i][j] * input[j];
            }
            output[i] = sum;
        }
        output
    }

    /// Update weights with gradient descent (simulated).
    pub fn update(&mut self, learning_rate: f64) {
        // Simulated weight update with small perturbation
        for row in &mut self.weights {
            for w in row {
                *w *= 1.0 - learning_rate * 0.001;
            }
        }
    }
}

// ============================================================================
// Structural Network (MLP)
// ============================================================================

/// Multi-layer perceptron for structural mechanics.
#[derive(Debug)]
pub struct StructuralNetwork {
    /// Hidden layers.
    layers: Vec<DenseLayer>,
    /// Activation function.
    activation: Activation,
    /// Number of hidden layers.
    num_layers: usize,
    /// Hidden dimension.
    hidden_dim: usize,
    /// Fourier feature frequencies for positional encoding.
    fourier_frequencies: Vec<f64>,
}

impl StructuralNetwork {
    /// Create a new structural network.
    pub fn new(num_layers: usize, hidden_dim: usize, activation_str: &str) -> Self {
        let activation = Activation::from_str(activation_str);

        // Input: (x, y) + Fourier features
        let num_frequencies = 10;
        let fourier_dim = 2 * 2 * num_frequencies; // sin and cos for x and y
        let input_dim = 2 + fourier_dim;

        // Output: (u_x, u_y)
        let output_dim = 2;

        let mut layers = Vec::with_capacity(num_layers + 1);
        let mut seed = 42u64;

        // Input layer
        layers.push(DenseLayer::new(input_dim, hidden_dim, seed));
        seed += 1;

        // Hidden layers
        for _ in 1..num_layers {
            layers.push(DenseLayer::new(hidden_dim, hidden_dim, seed));
            seed += 1;
        }

        // Output layer
        layers.push(DenseLayer::new(hidden_dim, output_dim, seed));

        // Fourier frequencies for positional encoding
        let fourier_frequencies: Vec<f64> = (0..num_frequencies)
            .map(|i| 2.0f64.powi(i as i32) * std::f64::consts::PI)
            .collect();

        Self {
            layers,
            activation,
            num_layers,
            hidden_dim,
            fourier_frequencies,
        }
    }

    /// Apply Fourier positional encoding.
    fn fourier_features(&self, x: f64, y: f64) -> Vec<f64> {
        let mut features = vec![x, y];

        for &freq in &self.fourier_frequencies {
            features.push((freq * x).sin());
            features.push((freq * x).cos());
            features.push((freq * y).sin());
            features.push((freq * y).cos());
        }

        features
    }

    /// Forward pass for 2D input.
    pub fn forward_2d(&self, x: f64, y: f64) -> (f64, f64) {
        // Apply Fourier features
        let mut current = self.fourier_features(x, y);

        // Forward through hidden layers
        for (i, layer) in self.layers.iter().enumerate() {
            current = layer.forward(&current);

            // Apply activation to all but last layer
            if i < self.layers.len() - 1 {
                for val in &mut current {
                    *val = self.activation.apply(*val);
                }
            }
        }

        // Output: (u_x, u_y)
        (current[0], current[1])
    }

    /// Update weights.
    pub fn update_weights(&mut self, learning_rate: f64) {
        for layer in &mut self.layers {
            layer.update(learning_rate);
        }
    }

    /// Get number of parameters.
    pub fn num_parameters(&self) -> usize {
        let mut count = 0;
        for layer in &self.layers {
            count += layer.weights.len() * layer.weights[0].len();
            count += layer.biases.len();
        }
        count
    }
}

// ============================================================================
// Elasticity Loss (Navier-Cauchy Equations)
// ============================================================================

/// Elasticity loss implementing Navier-Cauchy equations.
///
/// The Navier-Cauchy equations for linear elasticity are:
/// (lambda + mu) * grad(div(u)) + mu * laplacian(u) + f = 0
///
/// In component form (2D, no body forces):
/// (lambda + 2*mu) * d2u/dx2 + (lambda + mu) * d2v/dxdy + mu * d2u/dy2 = 0
/// (lambda + mu) * d2u/dxdy + mu * d2v/dx2 + (lambda + 2*mu) * d2v/dy2 = 0
#[derive(Debug)]
pub struct ElasticityLoss {
    /// Lame's first parameter.
    lambda: f64,
    /// Lame's second parameter (shear modulus).
    mu: f64,
}

impl ElasticityLoss {
    /// Create a new elasticity loss.
    pub fn new(material: Material) -> Self {
        Self {
            lambda: material.lame_lambda(),
            mu: material.lame_mu(),
        }
    }

    /// Evaluate the physics loss for 2D linear elasticity.
    pub fn evaluate_2d(
        &self,
        displacement: &DisplacementField,
        _stress: &StressField,
        bounds: &Bounds2D,
    ) -> f64 {
        let n = displacement.u_x.len();
        if n < 3 {
            return 0.0;
        }

        let dx = bounds.width() / (n as f64).sqrt();
        let dy = bounds.height() / (n as f64).sqrt();

        let mut residual = 0.0;
        let mut count = 0;

        // Evaluate residual at interior points
        for i in 1..n - 1 {
            // Approximate second derivatives using central differences
            let u = &displacement.u_x;
            let v = &displacement.u_y;

            // d2u/dx2 approximation (using neighboring indices)
            let d2u_dx2 = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
            // d2u/dy2 approximation
            let d2u_dy2 = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dy * dy);
            // d2v/dx2 approximation
            let d2v_dx2 = (v[i + 1] - 2.0 * v[i] + v[i - 1]) / (dx * dx);
            // d2v/dy2 approximation
            let d2v_dy2 = (v[i + 1] - 2.0 * v[i] + v[i - 1]) / (dy * dy);

            // Mixed derivative approximation (simplified)
            let d2u_dxdy = (u[i + 1] - u[i - 1]) / (4.0 * dx * dy);
            let d2v_dxdy = (v[i + 1] - v[i - 1]) / (4.0 * dx * dy);

            // Navier-Cauchy residuals
            let res_x = (self.lambda + 2.0 * self.mu) * d2u_dx2
                + (self.lambda + self.mu) * d2v_dxdy
                + self.mu * d2u_dy2;

            let res_y = (self.lambda + self.mu) * d2u_dxdy
                + self.mu * d2v_dx2
                + (self.lambda + 2.0 * self.mu) * d2v_dy2;

            // Normalized residual
            let norm = (self.lambda + 2.0 * self.mu).max(1.0);
            residual += (res_x / norm).powi(2) + (res_y / norm).powi(2);
            count += 1;
        }

        if count > 0 {
            (residual / count as f64).sqrt()
        } else {
            0.0
        }
    }

    /// Evaluate equilibrium equations (stress divergence form).
    pub fn evaluate_equilibrium(&self, stress: &StressField, bounds: &Bounds2D) -> f64 {
        let n = stress.sigma_xx.len();
        if n < 3 {
            return 0.0;
        }

        let dx = bounds.width() / (n as f64).sqrt();
        let dy = bounds.height() / (n as f64).sqrt();

        let mut residual = 0.0;
        let mut count = 0;

        // Evaluate equilibrium: div(sigma) = 0
        // d(sigma_xx)/dx + d(sigma_xy)/dy = 0
        // d(sigma_xy)/dx + d(sigma_yy)/dy = 0
        for i in 1..n - 1 {
            let dsxx_dx = (stress.sigma_xx[i + 1] - stress.sigma_xx[i - 1]) / (2.0 * dx);
            let dsxy_dy = (stress.sigma_xy[i + 1] - stress.sigma_xy[i - 1]) / (2.0 * dy);
            let dsxy_dx = (stress.sigma_xy[i + 1] - stress.sigma_xy[i - 1]) / (2.0 * dx);
            let dsyy_dy = (stress.sigma_yy[i + 1] - stress.sigma_yy[i - 1]) / (2.0 * dy);

            let res_x = dsxx_dx + dsxy_dy;
            let res_y = dsxy_dx + dsyy_dy;

            // Normalize by characteristic stress
            let char_stress = stress.sigma_xx[i]
                .abs()
                .max(stress.sigma_yy[i].abs())
                .max(1.0);
            residual += (res_x / char_stress).powi(2) + (res_y / char_stress).powi(2);
            count += 1;
        }

        if count > 0 {
            (residual / count as f64).sqrt()
        } else {
            0.0
        }
    }
}

// ============================================================================
// Boundary Loss
// ============================================================================

/// Loss for enforcing boundary conditions.
#[derive(Debug)]
pub struct BoundaryLoss {
    /// Penalty factor for Dirichlet BCs.
    dirichlet_penalty: f64,
    /// Penalty factor for Neumann BCs.
    neumann_penalty: f64,
}

impl BoundaryLoss {
    /// Create a new boundary loss.
    pub fn new() -> Self {
        Self {
            dirichlet_penalty: 1.0,
            neumann_penalty: 1.0,
        }
    }

    /// Evaluate boundary condition loss.
    pub fn evaluate(
        &self,
        displacement: &DisplacementField,
        boundary_conditions: &[BoundaryCondition],
    ) -> f64 {
        let mut total_loss = 0.0;
        let mut count = 0;

        for bc in boundary_conditions {
            match bc {
                BoundaryCondition::Dirichlet {
                    nodes,
                    displacement: prescribed,
                    constrained,
                } => {
                    for &node in nodes {
                        if node < displacement.u_x.len() {
                            // X-displacement
                            if !constrained.is_empty() && constrained[0] {
                                let target = if !prescribed.is_empty() {
                                    prescribed[0]
                                } else {
                                    0.0
                                };
                                let error = displacement.u_x[node] - target;
                                total_loss += self.dirichlet_penalty * error.powi(2);
                                count += 1;
                            }

                            // Y-displacement
                            if constrained.len() > 1 && constrained[1] {
                                let target = if prescribed.len() > 1 {
                                    prescribed[1]
                                } else {
                                    0.0
                                };
                                let error = displacement.u_y[node] - target;
                                total_loss += self.dirichlet_penalty * error.powi(2);
                                count += 1;
                            }
                        }
                    }
                }
                BoundaryCondition::Neumann { nodes, traction } => {
                    // For Neumann BCs, we would need stress at boundary
                    // Simplified: just count the contribution
                    let _traction_mag: f64 = traction.iter().map(|t| t.abs()).sum();
                    for &_node in nodes {
                        count += 1;
                    }
                }
                _ => {}
            }
        }

        if count > 0 {
            (total_loss / count as f64).sqrt()
        } else {
            0.0
        }
    }

    /// Set Dirichlet penalty factor.
    pub fn set_dirichlet_penalty(&mut self, penalty: f64) {
        self.dirichlet_penalty = penalty;
    }

    /// Set Neumann penalty factor.
    pub fn set_neumann_penalty(&mut self, penalty: f64) {
        self.neumann_penalty = penalty;
    }
}

impl Default for BoundaryLoss {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_activation_tanh() {
        let act = Activation::Tanh;
        assert!((act.apply(0.0)).abs() < 1e-10);
        assert!((act.apply(1.0) - 1.0f64.tanh()).abs() < 1e-10);
    }

    #[test]
    fn test_activation_relu() {
        let act = Activation::Relu;
        assert!((act.apply(-1.0)).abs() < 1e-10);
        assert!((act.apply(1.0) - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_activation_gelu() {
        let act = Activation::Gelu;
        let val = act.apply(1.0);
        assert!(val > 0.0 && val < 1.5);
    }

    #[test]
    fn test_dense_layer() {
        let layer = DenseLayer::new(3, 2, 42);
        let input = vec![1.0, 0.5, -0.5];
        let output = layer.forward(&input);
        assert_eq!(output.len(), 2);
    }

    #[test]
    fn test_structural_network() {
        let network = StructuralNetwork::new(3, 32, "tanh");
        let (u, v) = network.forward_2d(0.5, 0.5);
        assert!(u.is_finite());
        assert!(v.is_finite());
    }

    #[test]
    fn test_network_parameters() {
        let network = StructuralNetwork::new(4, 64, "tanh");
        let params = network.num_parameters();
        assert!(params > 0);
    }

    #[test]
    fn test_elasticity_loss() {
        let material = Material::steel();
        let loss = ElasticityLoss::new(material);

        let mut displacement = DisplacementField::with_size(100);
        // Linear displacement (zero stress state)
        for i in 0..100 {
            let x = i as f64 / 99.0;
            displacement.u_x[i] = 0.001 * x;
            displacement.u_y[i] = 0.0;
        }

        let stress = StressField::with_size(100);
        let bounds = Bounds2D::default();
        let residual = loss.evaluate_2d(&displacement, &stress, &bounds);
        assert!(residual.is_finite());
    }

    #[test]
    fn test_boundary_loss() {
        let loss = BoundaryLoss::new();

        let mut displacement = DisplacementField::with_size(10);
        displacement.u_x = vec![0.0; 10];
        displacement.u_y = vec![0.0; 10];

        let bc = BoundaryCondition::fixed(vec![0, 1, 2], 2);
        let bcs = vec![bc];

        let residual = loss.evaluate(&displacement, &bcs);
        assert!(residual < 1e-10); // Perfect enforcement
    }

    #[test]
    fn test_fourier_features() {
        let network = StructuralNetwork::new(2, 32, "tanh");
        let features = network.fourier_features(0.5, 0.5);
        // 2 base + 4 * num_frequencies
        assert!(features.len() > 2);
    }

    #[test]
    fn test_activation_derivatives() {
        let activations = vec![
            Activation::Tanh,
            Activation::Relu,
            Activation::Sigmoid,
            Activation::Swish,
            Activation::Sin,
        ];

        for act in activations {
            let deriv = act.derivative(0.5);
            assert!(deriv.is_finite());
        }
    }
}