//! 4D PINN network for temperature prediction T(x, y, z, t)
//!
//! Uses Learnable Fourier Feature Network (LFFN) architecture with
//! 4-dimensional input for spatial coordinates and time.

use crate::config::{Activation, NetworkConfig};
use bioheat_shared::Point3D;
use rand::Rng;
use serde::{Deserialize, Serialize};

/// 4D coordinate input to the network
#[derive(Debug, Clone, Copy)]
pub struct Coordinate4D {
    pub x: f64,
    pub y: f64,
    pub z: f64,
    pub t: f64,
}

impl Coordinate4D {
    /// Create new 4D coordinate
    #[must_use]
    pub const fn new(x: f64, y: f64, z: f64, t: f64) -> Self {
        Self { x, y, z, t }
    }

    /// Create from Point3D and time
    #[must_use]
    pub fn from_point_and_time(p: Point3D, t: f32) -> Self {
        Self::new(p.x as f64, p.y as f64, p.z as f64, t as f64)
    }

    /// Convert to array
    #[must_use]
    pub fn to_array(&self) -> [f64; 4] {
        [self.x, self.y, self.z, self.t]
    }
}

/// Thermal PINN network for predicting temperature T(x,y,z,t)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThermalPinn {
    /// Configuration
    config: NetworkConfig,
    /// Fourier feature matrix B [4 x fourier_features]
    /// Each row corresponds to frequencies for (x, y, z, t)
    fourier_b: Vec<Vec<f64>>,
    /// MLP weights for each layer
    weights: Vec<Vec<Vec<f64>>>,
    /// MLP biases for each layer
    biases: Vec<Vec<f64>>,
}

impl ThermalPinn {
    /// Create a new ThermalPinn with given configuration
    #[must_use]
    pub fn new(config: NetworkConfig) -> Self {
        let mut rng = rand::thread_rng();

        // Initialize Fourier feature matrix B
        // Shape: [4, fourier_features]
        let fourier_b: Vec<Vec<f64>> = (0..4)
            .map(|_| {
                (0..config.fourier_features)
                    .map(|_| rng.r#gen::<f64>() * config.fourier_scale as f64)
                    .collect()
            })
            .collect();

        // Build MLP layer sizes
        // Input: 2 * fourier_features (cos + sin features)
        let input_dim = 2 * config.fourier_features;
        let mut layer_sizes = vec![input_dim];
        layer_sizes.extend(config.hidden_layers.iter().copied());
        layer_sizes.push(1); // Output: temperature

        // Initialize weights and biases with Xavier initialization
        let mut weights = Vec::new();
        let mut biases = Vec::new();

        for i in 0..layer_sizes.len() - 1 {
            let fan_in = layer_sizes[i];
            let fan_out = layer_sizes[i + 1];
            let scale = (6.0 / (fan_in + fan_out) as f64).sqrt();

            // Weight matrix [fan_out x fan_in]
            let w: Vec<Vec<f64>> = (0..fan_out)
                .map(|_| (0..fan_in).map(|_| rng.gen_range(-scale..scale)).collect())
                .collect();
            weights.push(w);

            // Bias vector [fan_out]
            let b: Vec<f64> = (0..fan_out).map(|_| 0.0).collect();
            biases.push(b);
        }

        Self {
            config,
            fourier_b,
            weights,
            biases,
        }
    }

    /// Create with default configuration
    #[must_use]
    pub fn default_network() -> Self {
        Self::new(NetworkConfig::default())
    }

    /// Compute Fourier features for input coordinate
    fn fourier_features(&self, coord: &Coordinate4D) -> Vec<f64> {
        let input = coord.to_array();
        let mut features = Vec::with_capacity(2 * self.config.fourier_features);

        for j in 0..self.config.fourier_features {
            // Compute dot product: B[:, j] · input
            let dot: f64 = (0..4).map(|i| self.fourier_b[i][j] * input[i]).sum();
            let angle = 2.0 * std::f64::consts::PI * dot;

            // Append cos and sin features
            features.push(angle.cos());
            features.push(angle.sin());
        }

        features
    }

    /// Apply activation function
    fn activate(&self, x: f64) -> f64 {
        match self.config.activation {
            Activation::Tanh => x.tanh(),
            Activation::Swish => x * (1.0 / (1.0 + (-x).exp())), // x * sigmoid(x)
            Activation::Gelu => {
                // Approximate GELU
                0.5 * x * (1.0 + (0.7978845608 * (x + 0.044715 * x.powi(3))).tanh())
            }
            Activation::Sin => x.sin(),
        }
    }

    /// Forward pass through MLP
    fn mlp_forward(&self, features: &[f64]) -> f64 {
        let mut current = features.to_vec();

        for (layer_idx, (w, b)) in self.weights.iter().zip(&self.biases).enumerate() {
            let is_last = layer_idx == self.weights.len() - 1;

            let mut next = Vec::with_capacity(w.len());
            for (row, bias) in w.iter().zip(b) {
                let sum: f64 = row.iter().zip(&current).map(|(wi, xi)| wi * xi).sum();
                let activated = if is_last {
                    sum + bias // No activation on output layer
                } else {
                    self.activate(sum + bias)
                };
                next.push(activated);
            }
            current = next;
        }

        // Output is a scalar (temperature)
        current[0]
    }

    /// Predict temperature at a 4D coordinate
    #[must_use]
    pub fn forward(&self, coord: &Coordinate4D) -> f64 {
        let features = self.fourier_features(coord);
        self.mlp_forward(&features)
    }

    /// Predict temperature from components
    #[must_use]
    pub fn predict(&self, x: f64, y: f64, z: f64, t: f64) -> f64 {
        self.forward(&Coordinate4D::new(x, y, z, t))
    }

    /// Compute gradients using finite differences
    /// Returns (dT/dx, dT/dy, dT/dz, dT/dt)
    #[must_use]
    pub fn gradients(&self, coord: &Coordinate4D, eps: f64) -> (f64, f64, f64, f64) {
        let _t0 = self.forward(coord);

        // dT/dx
        let tx_plus = self.forward(&Coordinate4D::new(coord.x + eps, coord.y, coord.z, coord.t));
        let tx_minus = self.forward(&Coordinate4D::new(coord.x - eps, coord.y, coord.z, coord.t));
        let dt_dx = (tx_plus - tx_minus) / (2.0 * eps);

        // dT/dy
        let ty_plus = self.forward(&Coordinate4D::new(coord.x, coord.y + eps, coord.z, coord.t));
        let ty_minus = self.forward(&Coordinate4D::new(coord.x, coord.y - eps, coord.z, coord.t));
        let dt_dy = (ty_plus - ty_minus) / (2.0 * eps);

        // dT/dz
        let tz_plus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z + eps, coord.t));
        let tz_minus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z - eps, coord.t));
        let dt_dz = (tz_plus - tz_minus) / (2.0 * eps);

        // dT/dt
        let tt_plus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z, coord.t + eps));
        let tt_minus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z, coord.t - eps));
        let dt_dt = (tt_plus - tt_minus) / (2.0 * eps);

        (dt_dx, dt_dy, dt_dz, dt_dt)
    }

    /// Compute second derivatives (for Laplacian)
    /// Returns (d²T/dx², d²T/dy², d²T/dz²)
    #[must_use]
    pub fn second_derivatives(&self, coord: &Coordinate4D, eps: f64) -> (f64, f64, f64) {
        let t0 = self.forward(coord);

        // d²T/dx²
        let tx_plus = self.forward(&Coordinate4D::new(coord.x + eps, coord.y, coord.z, coord.t));
        let tx_minus = self.forward(&Coordinate4D::new(coord.x - eps, coord.y, coord.z, coord.t));
        let d2t_dx2 = (tx_plus - 2.0 * t0 + tx_minus) / (eps * eps);

        // d²T/dy²
        let ty_plus = self.forward(&Coordinate4D::new(coord.x, coord.y + eps, coord.z, coord.t));
        let ty_minus = self.forward(&Coordinate4D::new(coord.x, coord.y - eps, coord.z, coord.t));
        let d2t_dy2 = (ty_plus - 2.0 * t0 + ty_minus) / (eps * eps);

        // d²T/dz²
        let tz_plus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z + eps, coord.t));
        let tz_minus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z - eps, coord.t));
        let d2t_dz2 = (tz_plus - 2.0 * t0 + tz_minus) / (eps * eps);

        (d2t_dx2, d2t_dy2, d2t_dz2)
    }

    /// Compute Laplacian ∇²T = d²T/dx² + d²T/dy² + d²T/dz²
    #[must_use]
    pub fn laplacian(&self, coord: &Coordinate4D, eps: f64) -> f64 {
        let (d2x, d2y, d2z) = self.second_derivatives(coord, eps);
        d2x + d2y + d2z
    }

    /// Get number of trainable parameters
    #[must_use]
    pub fn num_parameters(&self) -> usize {
        let fourier_params = 4 * self.config.fourier_features;
        let mlp_params: usize = self
            .weights
            .iter()
            .zip(&self.biases)
            .map(|(w, b)| w.len() * w[0].len() + b.len())
            .sum();
        fourier_params + mlp_params
    }

    /// Get mutable access to weights for training
    pub fn weights_mut(&mut self) -> &mut Vec<Vec<Vec<f64>>> {
        &mut self.weights
    }

    /// Get mutable access to biases for training
    pub fn biases_mut(&mut self) -> &mut Vec<Vec<f64>> {
        &mut self.biases
    }

    /// Get mutable access to Fourier features for training
    pub fn fourier_b_mut(&mut self) -> &mut Vec<Vec<f64>> {
        &mut self.fourier_b
    }
}

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

    fn test_config() -> NetworkConfig {
        NetworkConfig {
            fourier_features: 16,
            fourier_scale: 2.0,
            hidden_layers: vec![32, 32],
            activation: Activation::Tanh,
        }
    }

    #[test]
    fn test_network_creation() {
        let net = ThermalPinn::new(test_config());
        assert!(net.num_parameters() > 0);
    }

    #[test]
    fn test_forward_pass() {
        let net = ThermalPinn::new(test_config());
        let coord = Coordinate4D::new(0.0, 0.0, 0.0, 0.0);
        let t = net.forward(&coord);
        // Output should be a finite number
        assert!(t.is_finite());
    }

    #[test]
    fn test_gradients() {
        let net = ThermalPinn::new(test_config());
        let coord = Coordinate4D::new(0.05, 0.05, 0.05, 100.0);
        let (dx, dy, dz, dt) = net.gradients(&coord, 1e-5);

        // Gradients should be finite
        assert!(dx.is_finite());
        assert!(dy.is_finite());
        assert!(dz.is_finite());
        assert!(dt.is_finite());
    }

    #[test]
    fn test_laplacian() {
        let net = ThermalPinn::new(test_config());
        let coord = Coordinate4D::new(0.05, 0.05, 0.05, 100.0);
        let lap = net.laplacian(&coord, 1e-4);
        assert!(lap.is_finite());
    }

    #[test]
    fn test_batch_consistency() {
        let net = ThermalPinn::new(test_config());
        let coord = Coordinate4D::new(0.03, 0.04, 0.02, 50.0);

        // Forward pass should be deterministic
        let t1 = net.forward(&coord);
        let t2 = net.forward(&coord);
        assert!((t1 - t2).abs() < 1e-10);
    }
}