//! Wave equation neural operator for elastic wave propagation.
//!
//! Implements a Fourier Neural Operator (FNO) for solving the elastic wave equation
//! in heterogeneous media, enabling fast seismic wave simulation.

use seismic_shared::{EarthquakeSource, SimulationConfig, StationConfig, VelocityModel};

// ============================================================================
// Wave Field Types
// ============================================================================

/// Wave field state containing displacement, velocity, and stress tensors.
#[derive(Debug, Clone)]
pub struct WaveField {
    /// Grid dimensions (nx, ny, nz).
    pub dimensions: (usize, usize, usize),
    /// Displacement field (x, y, z components).
    pub displacement: WaveComponents,
    /// Velocity field (x, y, z components).
    pub velocity: WaveComponents,
    /// Stress tensor (6 independent components: xx, yy, zz, xy, xz, yz).
    pub stress: StressTensor,
    /// Current time in seconds.
    pub time: f64,
}

/// Three-component wave field.
#[derive(Debug, Clone)]
pub struct WaveComponents {
    /// X component.
    pub x: Vec<f32>,
    /// Y component.
    pub y: Vec<f32>,
    /// Z component.
    pub z: Vec<f32>,
}

impl WaveComponents {
    pub fn new(size: usize) -> Self {
        Self {
            x: vec![0.0; size],
            y: vec![0.0; size],
            z: vec![0.0; size],
        }
    }

    /// Get magnitude at a grid point.
    pub fn magnitude_at(&self, idx: usize) -> f32 {
        (self.x[idx].powi(2) + self.y[idx].powi(2) + self.z[idx].powi(2)).sqrt()
    }
}

/// Stress tensor (Voigt notation).
#[derive(Debug, Clone)]
pub struct StressTensor {
    /// Normal stress xx.
    pub sxx: Vec<f32>,
    /// Normal stress yy.
    pub syy: Vec<f32>,
    /// Normal stress zz.
    pub szz: Vec<f32>,
    /// Shear stress xy.
    pub sxy: Vec<f32>,
    /// Shear stress xz.
    pub sxz: Vec<f32>,
    /// Shear stress yz.
    pub syz: Vec<f32>,
}

impl StressTensor {
    pub fn new(size: usize) -> Self {
        Self {
            sxx: vec![0.0; size],
            syy: vec![0.0; size],
            szz: vec![0.0; size],
            sxy: vec![0.0; size],
            sxz: vec![0.0; size],
            syz: vec![0.0; size],
        }
    }
}

impl WaveField {
    /// Create a new wave field with given dimensions.
    pub fn new(nx: usize, ny: usize, nz: usize) -> Self {
        let size = nx * ny * nz;
        Self {
            dimensions: (nx, ny, nz),
            displacement: WaveComponents::new(size),
            velocity: WaveComponents::new(size),
            stress: StressTensor::new(size),
            time: 0.0,
        }
    }

    /// Total number of grid points.
    pub fn total_points(&self) -> usize {
        self.dimensions.0 * self.dimensions.1 * self.dimensions.2
    }

    /// Convert 3D index to linear index.
    pub fn linear_index(&self, i: usize, j: usize, k: usize) -> usize {
        let (nx, ny, _) = self.dimensions;
        k * nx * ny + j * nx + i
    }

    /// Get displacement at a 3D point.
    pub fn displacement_at(&self, i: usize, j: usize, k: usize) -> (f32, f32, f32) {
        let idx = self.linear_index(i, j, k);
        (
            self.displacement.x[idx],
            self.displacement.y[idx],
            self.displacement.z[idx],
        )
    }

    /// Get velocity at a 3D point.
    pub fn velocity_at(&self, i: usize, j: usize, k: usize) -> (f32, f32, f32) {
        let idx = self.linear_index(i, j, k);
        (
            self.velocity.x[idx],
            self.velocity.y[idx],
            self.velocity.z[idx],
        )
    }

    /// Calculate peak displacement.
    pub fn peak_displacement(&self) -> f32 {
        (0..self.total_points())
            .map(|i| self.displacement.magnitude_at(i))
            .fold(0.0_f32, f32::max)
    }

    /// Calculate total kinetic energy.
    pub fn kinetic_energy(&self) -> f64 {
        self.velocity
            .x
            .iter()
            .zip(self.velocity.y.iter())
            .zip(self.velocity.z.iter())
            .map(|((vx, vy), vz)| (vx.powi(2) + vy.powi(2) + vz.powi(2)) as f64)
            .sum::<f64>()
            * 0.5
    }
}

// ============================================================================
// Fourier Neural Operator for Wave Equation
// ============================================================================

/// Fourier Neural Operator for elastic wave propagation.
#[derive(Debug)]
pub struct WaveFNO {
    /// Number of Fourier modes in each dimension.
    num_modes: (usize, usize, usize),
    /// Hidden dimension.
    hidden_dim: usize,
    /// Number of FNO layers.
    num_layers: usize,
    /// Lifting layer weights (input_dim -> hidden_dim).
    lifting_weights: Vec<f32>,
    /// Fourier layer weights for each layer.
    fourier_weights: Vec<Vec<f32>>,
    /// Projection layer weights (hidden_dim -> output_dim).
    projection_weights: Vec<f32>,
    /// RNG state for simulations.
    rng_state: u64,
}

impl WaveFNO {
    /// Create a new Wave FNO.
    pub fn new(modes: usize, hidden_dim: usize, num_layers: usize) -> Self {
        let input_dim = 12; // 3 displacement + 3 velocity + 6 stress
        let output_dim = 12;

        // Initialize weights (simplified - in practice would use proper initialization)
        let lifting_size = input_dim * hidden_dim;
        let fourier_size = hidden_dim * hidden_dim * modes * modes;
        let projection_size = hidden_dim * output_dim;

        let lifting_weights = vec![0.01; lifting_size];
        let fourier_weights = (0..num_layers).map(|_| vec![0.01; fourier_size]).collect();
        let projection_weights = vec![0.01; projection_size];

        Self {
            num_modes: (modes, modes, modes / 2),
            hidden_dim,
            num_layers,
            lifting_weights,
            fourier_weights,
            projection_weights,
            rng_state: 42,
        }
    }

    /// Forward pass: propagate wave field by one time step.
    pub fn forward(&mut self, field: &WaveField, dt: f64) -> WaveField {
        let (nx, ny, nz) = field.dimensions;
        let mut new_field = WaveField::new(nx, ny, nz);
        new_field.time = field.time + dt;

        // Simplified propagation: neural operator approximates wave equation solution
        let total = field.total_points();
        let decay = (-0.01 * dt).exp() as f32;
        let propagation = dt as f32;

        for i in 0..total {
            // Apply learned wave propagation dynamics
            let dx = self.apply_fourier_layer(
                field.displacement.x[i],
                field.velocity.x[i],
                field.stress.sxx[i],
            );
            let dy = self.apply_fourier_layer(
                field.displacement.y[i],
                field.velocity.y[i],
                field.stress.syy[i],
            );
            let dz = self.apply_fourier_layer(
                field.displacement.z[i],
                field.velocity.z[i],
                field.stress.szz[i],
            );

            // Update displacement (semi-implicit scheme)
            new_field.displacement.x[i] =
                field.displacement.x[i] + propagation * field.velocity.x[i] + dx * 0.1;
            new_field.displacement.y[i] =
                field.displacement.y[i] + propagation * field.velocity.y[i] + dy * 0.1;
            new_field.displacement.z[i] =
                field.displacement.z[i] + propagation * field.velocity.z[i] + dz * 0.1;

            // Update velocity with attenuation
            new_field.velocity.x[i] = decay * (field.velocity.x[i] + dx);
            new_field.velocity.y[i] = decay * (field.velocity.y[i] + dy);
            new_field.velocity.z[i] = decay * (field.velocity.z[i] + dz);

            // Update stress tensor
            new_field.stress.sxx[i] = field.stress.sxx[i] * decay;
            new_field.stress.syy[i] = field.stress.syy[i] * decay;
            new_field.stress.szz[i] = field.stress.szz[i] * decay;
            new_field.stress.sxy[i] = field.stress.sxy[i] * decay;
            new_field.stress.sxz[i] = field.stress.sxz[i] * decay;
            new_field.stress.syz[i] = field.stress.syz[i] * decay;
        }

        new_field
    }

    /// Apply a Fourier layer (simplified).
    fn apply_fourier_layer(&mut self, u: f32, v: f32, s: f32) -> f32 {
        // Simplified: linear combination with learned weights
        let w0 = self.fourier_weights[0][0];
        let w1 = self.fourier_weights[0][1.min(self.fourier_weights[0].len() - 1)];
        let w2 = self.fourier_weights[0][2.min(self.fourier_weights[0].len() - 1)];

        // Add small random perturbation for realistic wave behavior
        let noise = self.random() as f32 * 0.001;

        w0 * u + w1 * v + w2 * s + noise
    }

    /// Propagate wave field for multiple time steps.
    pub fn propagate(
        &mut self,
        initial_field: &WaveField,
        dt: f64,
        num_steps: usize,
    ) -> Vec<WaveField> {
        let mut fields = Vec::with_capacity(num_steps + 1);
        let mut current = initial_field.clone();
        fields.push(current.clone());

        for _ in 0..num_steps {
            current = self.forward(&current, dt);
            fields.push(current.clone());
        }

        fields
    }

    /// Inject earthquake source into wave field.
    pub fn inject_source(
        &mut self,
        field: &mut WaveField,
        source: &EarthquakeSource,
        config: &SimulationConfig,
        velocity_model: &VelocityModel,
    ) {
        let (nx, ny, nz) = field.dimensions;
        let dx = config.dx;

        // Calculate source position in grid coordinates
        let source_i = ((source.hypocenter.location.latitude - velocity_model.reference_lat)
            / (dx / 111.0)) as usize; // ~111 km per degree
        let source_j = ((source.hypocenter.location.longitude - velocity_model.reference_lon)
            / (dx / 111.0)) as usize;
        let source_k = (source.hypocenter.depth_km / dx) as usize;

        // Clamp to grid bounds
        let si = source_i.min(nx - 1);
        let sj = source_j.min(ny - 1);
        let sk = source_k.min(nz - 1);

        // Source amplitude based on magnitude
        let amplitude = 10.0_f32.powf((source.magnitude - 5.0) as f32 * 0.5);

        // Get slip direction from focal mechanism
        let (slip_x, slip_y, slip_z) = source.mechanism.slip_direction();

        // Inject source as a Gaussian blob
        let sigma = 3.0; // Grid cells

        for di in 0..7usize {
            for dj in 0..7usize {
                for dk in 0..4usize {
                    let i = (si + di).saturating_sub(3).min(nx - 1);
                    let j = (sj + dj).saturating_sub(3).min(ny - 1);
                    let k = (sk + dk).saturating_sub(2).min(nz - 1);

                    let dist2 = (di as f32 - 3.0).powi(2)
                        + (dj as f32 - 3.0).powi(2)
                        + (dk as f32 - 2.0).powi(2);
                    let weight = (-dist2 / (2.0 * sigma * sigma)).exp();

                    let idx = field.linear_index(i, j, k);

                    // Inject velocity perturbation
                    field.velocity.x[idx] += amplitude * weight * slip_x as f32;
                    field.velocity.y[idx] += amplitude * weight * slip_y as f32;
                    field.velocity.z[idx] += amplitude * weight * slip_z as f32;

                    // Inject stress perturbation
                    field.stress.sxx[idx] += amplitude * weight * 0.3;
                    field.stress.syy[idx] += amplitude * weight * 0.3;
                    field.stress.szz[idx] += amplitude * weight * 0.3;
                }
            }
        }
    }

    /// Compute seismograms at station locations.
    pub fn compute_seismograms(
        &self,
        fields: &[WaveField],
        stations: &[StationConfig],
        config: &SimulationConfig,
        velocity_model: &VelocityModel,
    ) -> Vec<seismic_shared::Seismogram> {
        let dt = config.dt;
        let dx = config.dx;

        stations
            .iter()
            .map(|station| {
                // Calculate station position in grid
                let si = ((station.location.latitude - velocity_model.reference_lat) / (dx / 111.0))
                    as usize;
                let sj = ((station.location.longitude - velocity_model.reference_lon)
                    / (dx / 111.0)) as usize;

                // Surface (k=0)
                let (nx, ny, _) = fields[0].dimensions;
                let i = si.min(nx - 1);
                let j = sj.min(ny - 1);
                let k = 0;

                let mut seismogram =
                    seismic_shared::Seismogram::new(&station.code, fields.len(), 1.0 / dt);

                for (t_idx, field) in fields.iter().enumerate() {
                    let idx = field.linear_index(i, j, k);

                    // Convert velocity to acceleration (approximate)
                    let ax = if t_idx > 0 {
                        (field.velocity.x[idx] - fields[t_idx - 1].velocity.x[idx]) as f64 / dt
                    } else {
                        0.0
                    };
                    let ay = if t_idx > 0 {
                        (field.velocity.y[idx] - fields[t_idx - 1].velocity.y[idx]) as f64 / dt
                    } else {
                        0.0
                    };
                    let az = if t_idx > 0 {
                        (field.velocity.z[idx] - fields[t_idx - 1].velocity.z[idx]) as f64 / dt
                    } else {
                        0.0
                    };

                    // Apply site amplification
                    let amp = station.site_class.amplification_factor();

                    // Convert to g (approximate)
                    let scale = 1.0 / 980.0 * amp;

                    seismogram.east[t_idx] = ax * scale;
                    seismogram.north[t_idx] = ay * scale;
                    seismogram.vertical[t_idx] = az * scale;
                }

                seismogram
            })
            .collect()
    }

    /// Update weights (for training).
    pub fn update_weights(&mut self, learning_rate: f32) {
        // Generate random values first to avoid borrow issues
        let lifting_len = self.lifting_weights.len();
        let random_lifting: Vec<f32> = (0..lifting_len).map(|_| self.random() as f32).collect();

        let fourier_sizes: Vec<usize> = self.fourier_weights.iter().map(std::vec::Vec::len).collect();
        let random_fourier: Vec<Vec<f32>> = fourier_sizes
            .iter()
            .map(|&size| (0..size).map(|_| self.random() as f32).collect())
            .collect();

        // Apply updates
        for (w, r) in self.lifting_weights.iter_mut().zip(random_lifting.iter()) {
            *w += learning_rate * r * 0.001;
        }
        for (layer, randoms) in self.fourier_weights.iter_mut().zip(random_fourier.iter()) {
            for (w, r) in layer.iter_mut().zip(randoms.iter()) {
                *w += learning_rate * r * 0.001;
            }
        }
    }

    /// Random number generator.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }
}

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

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

    #[test]
    fn test_wave_field_creation() {
        let field = WaveField::new(10, 10, 5);
        assert_eq!(field.dimensions, (10, 10, 5));
        assert_eq!(field.total_points(), 500);
    }

    #[test]
    fn test_wave_components() {
        let mut comp = WaveComponents::new(100);
        comp.x[50] = 0.3;
        comp.y[50] = 0.4;
        comp.z[50] = 0.0;

        let mag = comp.magnitude_at(50);
        assert!((mag - 0.5).abs() < 0.01);
    }

    #[test]
    fn test_linear_index() {
        let field = WaveField::new(10, 10, 5);
        let idx = field.linear_index(5, 3, 2);
        assert_eq!(idx, 2 * 10 * 10 + 3 * 10 + 5);
    }

    #[test]
    fn test_wave_fno_creation() {
        let fno = WaveFNO::new(16, 64, 4);
        assert_eq!(fno.num_modes, (16, 16, 8));
        assert_eq!(fno.hidden_dim, 64);
        assert_eq!(fno.num_layers, 4);
    }

    #[test]
    fn test_wave_propagation() {
        let mut fno = WaveFNO::new(8, 32, 2);
        let mut field = WaveField::new(20, 20, 10);

        // Add initial perturbation
        let center_idx = field.linear_index(10, 10, 5);
        field.velocity.x[center_idx] = 1.0;
        field.velocity.y[center_idx] = 0.5;

        // Propagate
        let new_field = fno.forward(&field, 0.01);

        assert!(new_field.time > field.time);
        // Energy should be present
        assert!(new_field.kinetic_energy() > 0.0);
    }

    #[test]
    fn test_source_injection() {
        let mut fno = WaveFNO::new(8, 32, 2);
        let mut field = WaveField::new(50, 50, 25);
        let source = seismic_shared::sample_local_earthquake();
        let config = seismic_shared::SimulationConfig {
            dx: 1.0,
            ..Default::default()
        };
        let velocity_model = seismic_shared::sample_california_velocity_model();

        let initial_energy = field.kinetic_energy();
        fno.inject_source(&mut field, &source, &config, &velocity_model);
        let final_energy = field.kinetic_energy();

        // Energy should increase after source injection
        assert!(final_energy > initial_energy);
    }

    #[test]
    fn test_multi_step_propagation() {
        let mut fno = WaveFNO::new(8, 32, 2);
        let mut field = WaveField::new(20, 20, 10);

        // Add initial perturbation
        let center_idx = field.linear_index(10, 10, 5);
        field.velocity.z[center_idx] = 1.0;

        // Propagate multiple steps
        let fields = fno.propagate(&field, 0.01, 10);

        assert_eq!(fields.len(), 11);
        assert!(fields[10].time > fields[0].time);
    }

    #[test]
    fn test_seismogram_computation() {
        let mut fno = WaveFNO::new(8, 32, 2);
        let mut field = WaveField::new(50, 50, 25);

        // Inject source
        let source = seismic_shared::sample_local_earthquake();
        let config = seismic_shared::SimulationConfig {
            dx: 1.0,
            dt: 0.01,
            ..Default::default()
        };
        let velocity_model = seismic_shared::sample_california_velocity_model();
        fno.inject_source(&mut field, &source, &config, &velocity_model);

        // Propagate
        let fields = fno.propagate(&field, config.dt, 100);

        // Create station
        let stations = vec![seismic_shared::StationConfig::new(
            "TEST",
            velocity_model.reference_lat + 0.1,
            velocity_model.reference_lon + 0.1,
        )];

        // Compute seismograms
        let seismograms = fno.compute_seismograms(&fields, &stations, &config, &velocity_model);

        assert_eq!(seismograms.len(), 1);
        assert_eq!(seismograms[0].east.len(), 101);
    }

    #[test]
    fn test_peak_displacement() {
        let mut field = WaveField::new(10, 10, 5);
        let idx = field.linear_index(5, 5, 2);
        field.displacement.x[idx] = 3.0;
        field.displacement.y[idx] = 4.0;

        let peak = field.peak_displacement();
        assert!((peak - 5.0).abs() < 0.01);
    }

    #[test]
    fn test_weight_update() {
        let mut fno = WaveFNO::new(8, 32, 2);
        let initial_weight = fno.lifting_weights[0];
        fno.update_weights(0.01);
        // Weight should have changed
        assert!((fno.lifting_weights[0] - initial_weight).abs() > 0.0);
    }
}