//! SeismicAI: Earthquake Simulation & Early Warning with Neural Operators
//!
//! This demo showcases neural operator-based seismic wave propagation for
//! earthquake simulation and early warning systems. It demonstrates:
//! - Fourier Neural Operators for elastic wave propagation
//! - Ground motion prediction equations (GMPEs)
//! - Real-time earthquake early warning
//! - Site effects and attenuation modeling

pub mod attenuation;
pub mod sample_data;
pub mod wave_operator;

use attenuation::{AttenuationModel, SiteEffects, GMPE};
use seismic_shared::{
    AlertLevel, EarthquakeSource, EarthquakeWarning, GeoLocation, GroundMotion, SimulationConfig,
    SimulationResult, StationConfig, VelocityModel, WarningConfig, WaveFieldSnapshot,
};
use wave_operator::{WaveFNO, WaveField};

// ============================================================================
// SeismicAI System
// ============================================================================

/// Main SeismicAI system for earthquake simulation and early warning.
#[derive(Debug)]
pub struct SeismicAI {
    /// Wave operator for full waveform simulation.
    wave_operator: WaveFNO,
    /// Source encoder for earthquake parametrization.
    source_encoder: SourceEncoder,
    /// Arrival predictor for travel time estimation.
    arrival_predictor: ArrivalPredictor,
    /// Attenuation model for ground motion prediction.
    attenuation_model: AttenuationModel,
    /// Simulation configuration.
    config: SimulationConfig,
    /// Warning configuration.
    warning_config: WarningConfig,
    /// RNG state for simulations.
    rng_state: u64,
}

/// Encodes earthquake source parameters for neural network input.
#[derive(Debug)]
pub struct SourceEncoder {
    /// Hidden dimension.
    hidden_dim: usize,
    /// Encoding weights.
    weights: Vec<f32>,
}

impl SourceEncoder {
    pub fn new(hidden_dim: usize) -> Self {
        Self {
            hidden_dim,
            weights: vec![0.01; hidden_dim * 10], // 10 source parameters
        }
    }

    /// Encode earthquake source to feature vector.
    pub fn encode(&self, source: &EarthquakeSource) -> Vec<f32> {
        let mut features = Vec::with_capacity(self.hidden_dim);

        // Extract key source parameters
        let lat = source.hypocenter.location.latitude as f32;
        let lon = source.hypocenter.location.longitude as f32;
        let depth = source.hypocenter.depth_km as f32;
        let mag = source.magnitude as f32;
        let strike = source.mechanism.strike as f32 / 360.0;
        let dip = source.mechanism.dip as f32 / 90.0;
        let rake = (source.mechanism.rake as f32 + 180.0) / 360.0;

        // Simple linear encoding
        for i in 0..self.hidden_dim {
            let w_idx = i * 7;
            let feature = self.weights[w_idx] * lat
                + self.weights[w_idx + 1] * lon
                + self.weights[w_idx + 2] * depth
                + self.weights[w_idx + 3] * mag
                + self.weights[w_idx + 4] * strike
                + self.weights[w_idx + 5] * dip
                + self.weights[w_idx + 6] * rake;

            features.push(feature.tanh());
        }

        features
    }
}

/// Predicts P and S wave arrival times.
#[derive(Debug)]
pub struct ArrivalPredictor {
    /// Average P-wave velocity (km/s).
    avg_vp: f64,
    /// Average S-wave velocity (km/s).
    avg_vs: f64,
}

impl ArrivalPredictor {
    pub fn new() -> Self {
        Self {
            avg_vp: 6.0,
            avg_vs: 3.5,
        }
    }

    /// Update velocities from a velocity model.
    pub fn from_velocity_model(model: &VelocityModel) -> Self {
        let n = model.layers.len() as f64;
        let avg_vp = model.layers.iter().map(|l| l.vp).sum::<f64>() / n;
        let avg_vs = model.layers.iter().map(|l| l.vs).sum::<f64>() / n;
        Self { avg_vp, avg_vs }
    }

    /// Predict P-wave arrival time.
    pub fn predict_p_arrival(&self, distance_km: f64, depth_km: f64) -> f64 {
        let path = (distance_km.powi(2) + depth_km.powi(2)).sqrt();
        path / self.avg_vp
    }

    /// Predict S-wave arrival time.
    pub fn predict_s_arrival(&self, distance_km: f64, depth_km: f64) -> f64 {
        let path = (distance_km.powi(2) + depth_km.powi(2)).sqrt();
        path / self.avg_vs
    }

    /// Predict S-P time (useful for distance estimation).
    pub fn predict_sp_time(&self, distance_km: f64, depth_km: f64) -> f64 {
        self.predict_s_arrival(distance_km, depth_km)
            - self.predict_p_arrival(distance_km, depth_km)
    }

    /// Estimate distance from S-P time.
    pub fn estimate_distance(&self, sp_time: f64) -> f64 {
        // Simplified: assumes horizontal distance >> depth
        sp_time / (1.0 / self.avg_vs - 1.0 / self.avg_vp)
    }
}

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

impl SeismicAI {
    /// Create a new SeismicAI system.
    pub fn new(config: SimulationConfig) -> Self {
        Self {
            wave_operator: WaveFNO::new(config.fno_modes, config.hidden_dim, 4),
            source_encoder: SourceEncoder::new(config.hidden_dim),
            arrival_predictor: ArrivalPredictor::new(),
            attenuation_model: AttenuationModel::nga_west2(),
            config,
            warning_config: WarningConfig::default(),
            rng_state: 42,
        }
    }

    /// Create with specific warning configuration.
    pub fn with_warning_config(mut self, warning_config: WarningConfig) -> Self {
        self.warning_config = warning_config;
        self
    }

    /// Set the attenuation model.
    pub fn with_attenuation_model(mut self, model: AttenuationModel) -> Self {
        self.attenuation_model = model;
        self
    }

    /// Simulate seismic wave propagation.
    pub fn simulate(
        &mut self,
        source: &EarthquakeSource,
        stations: &[StationConfig],
        velocity_model: &VelocityModel,
    ) -> SimulationResult {
        let start = std::time::Instant::now();

        // Update arrival predictor with velocity model
        self.arrival_predictor = ArrivalPredictor::from_velocity_model(velocity_model);

        // Calculate grid dimensions
        let (dx, dy, dz) = (self.config.dx, self.config.dx, self.config.dx);
        let (lx, ly, lz) = self.config.domain_size;
        let nx = (lx / dx) as usize;
        let ny = (ly / dy) as usize;
        let nz = (lz / dz) as usize;

        // Initialize wave field
        let mut wave_field = WaveField::new(nx, ny, nz);

        // Inject source
        self.wave_operator
            .inject_source(&mut wave_field, source, &self.config, velocity_model);

        // Propagate waves
        let num_steps = (self.config.duration / self.config.dt) as usize;
        let fields = self
            .wave_operator
            .propagate(&wave_field, self.config.dt, num_steps);

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

        // Predict ground motion at each station
        let ground_motions = self.predict_ground_motion(source, stations, velocity_model);

        // Generate warning
        let warning = self.generate_warning(source, stations, velocity_model);

        // Create snapshots (sample every N steps)
        let snapshot_interval = num_steps / 10;
        let snapshots: Vec<WaveFieldSnapshot> = fields
            .iter()
            .enumerate()
            .filter(|(i, _)| i % snapshot_interval == 0)
            .take(10)
            .map(|(_, field)| self.create_snapshot(field))
            .collect();

        let computation_time_ms = start.elapsed().as_secs_f64() * 1000.0;

        SimulationResult {
            source: source.clone(),
            ground_motions,
            seismograms: Some(seismograms),
            snapshots: Some(snapshots),
            warning: Some(warning),
            computation_time_ms,
        }
    }

    /// Predict ground motion at stations using GMPE.
    pub fn predict_ground_motion(
        &self,
        source: &EarthquakeSource,
        stations: &[StationConfig],
        velocity_model: &VelocityModel,
    ) -> Vec<(StationConfig, GroundMotion)> {
        stations
            .iter()
            .map(|station| {
                // Get base prediction from GMPE
                let mut gm = self
                    .attenuation_model
                    .predict(source, station, velocity_model);

                // Apply site effects
                if self.config.include_site_effects {
                    let site_effects = SiteEffects::from_station(station);
                    gm = site_effects.apply(&gm);
                }

                (station.clone(), gm)
            })
            .collect()
    }

    /// Generate earthquake early warning.
    pub fn generate_warning(
        &mut self,
        source: &EarthquakeSource,
        stations: &[StationConfig],
        velocity_model: &VelocityModel,
    ) -> EarthquakeWarning {
        // Find stations that would detect the event
        let detecting_stations: Vec<_> = stations
            .iter()
            .filter(|s| {
                let dist = source.hypocenter.location.distance_km(&s.location);
                dist < self.warning_config.alert_radius_km
            })
            .collect();

        if detecting_stations.len() < self.warning_config.min_stations
            || source.magnitude < self.warning_config.min_magnitude
        {
            return EarthquakeWarning::default();
        }

        // Estimate magnitude (with uncertainty)
        let estimated_magnitude = source.magnitude + (self.random() - 0.5) * 0.3;

        // Estimate location (with uncertainty)
        let estimated_location = GeoLocation::new(
            source.hypocenter.location.latitude + (self.random() - 0.5) * 0.1,
            source.hypocenter.location.longitude + (self.random() - 0.5) * 0.1,
        );

        let estimated_depth = source.hypocenter.depth_km + (self.random() - 0.5) * 5.0;

        // Calculate warning time for a typical target location
        let target_distance = 50.0; // km
        let p_arrival = self
            .arrival_predictor
            .predict_p_arrival(target_distance, source.hypocenter.depth_km);
        let s_arrival = self
            .arrival_predictor
            .predict_s_arrival(target_distance, source.hypocenter.depth_km);

        // Time to shaking at target
        let processing_delay = 3.0; // seconds for detection and processing
        let time_to_shaking = (s_arrival - p_arrival - processing_delay).max(0.0);

        // Predict ground motion at target
        let target_station = StationConfig::new(
            "TARGET",
            source.hypocenter.location.latitude + 0.45,
            source.hypocenter.location.longitude,
        );
        let expected_gm = self
            .attenuation_model
            .predict(source, &target_station, velocity_model);

        // Determine alert level
        let alert_level = if expected_gm.mmi >= 7.0 {
            AlertLevel::Severe
        } else if expected_gm.mmi >= 5.0 {
            AlertLevel::Warning
        } else if expected_gm.mmi >= 4.0 {
            AlertLevel::Watch
        } else if expected_gm.mmi >= 3.0 {
            AlertLevel::Advisory
        } else {
            AlertLevel::None
        };

        // Calculate confidence based on number of detecting stations
        let confidence = (detecting_stations.len() as f64 / 10.0).min(0.95);

        EarthquakeWarning {
            event_id: source.event_id.clone(),
            alert_level,
            estimated_magnitude,
            estimated_location,
            estimated_depth_km: estimated_depth,
            detecting_stations: detecting_stations.len(),
            time_to_shaking,
            expected_ground_motion: expected_gm,
            warning_time: p_arrival + processing_delay,
            origin_time: 0.0,
            confidence,
        }
    }

    /// Create a wave field snapshot for visualization.
    fn create_snapshot(&self, field: &WaveField) -> WaveFieldSnapshot {
        let (nx, ny, _) = field.dimensions;

        // Extract surface (z=0) displacement
        let mut disp_x = vec![vec![0.0_f32; ny]; nx];
        let mut disp_y = vec![vec![0.0_f32; ny]; nx];
        let mut disp_z = vec![vec![0.0_f32; ny]; nx];

        for i in 0..nx {
            for j in 0..ny {
                let (dx, dy, dz) = field.displacement_at(i, j, 0);
                disp_x[i][j] = dx;
                disp_y[i][j] = dy;
                disp_z[i][j] = dz;
            }
        }

        // Grid coordinates
        let grid_x: Vec<f32> = (0..nx).map(|i| i as f32 * self.config.dx as f32).collect();
        let grid_y: Vec<f32> = (0..ny).map(|j| j as f32 * self.config.dx as f32).collect();

        WaveFieldSnapshot {
            time: field.time,
            displacement_x: disp_x,
            displacement_y: disp_y,
            displacement_z: disp_z,
            grid_x,
            grid_y,
        }
    }

    /// 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
    }
}

// ============================================================================
// Demo Entry Point
// ============================================================================

/// Run the SeismicAI demo.
pub fn run_demo() -> SimulationResult {
    // Create demo scenario
    let (source, stations, velocity_model, config) = sample_data::local_test_scenario();

    // Create SeismicAI system
    let mut seismic_ai = SeismicAI::new(config)
        .with_warning_config(sample_data::urban_warning_config())
        .with_attenuation_model(AttenuationModel::california());

    // Run simulation
    seismic_ai.simulate(&source, &stations, &velocity_model)
}

/// Run a major earthquake scenario.
pub fn run_major_scenario() -> SimulationResult {
    let (source, stations, velocity_model, config) = sample_data::hayward_scenario();

    let mut seismic_ai =
        SeismicAI::new(config).with_warning_config(sample_data::urban_warning_config());

    seismic_ai.simulate(&source, &stations, &velocity_model)
}

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

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

    #[test]
    fn test_seismic_ai_creation() {
        let config = SimulationConfig::default();
        let seismic_ai = SeismicAI::new(config);
        assert_eq!(seismic_ai.arrival_predictor.avg_vp, 6.0);
    }

    #[test]
    fn test_source_encoder() {
        let encoder = SourceEncoder::new(32);
        let source = sample_data::local_earthquake();
        let features = encoder.encode(&source);

        assert_eq!(features.len(), 32);
        // All features should be in tanh range [-1, 1]
        for f in &features {
            assert!(*f >= -1.0 && *f <= 1.0);
        }
    }

    #[test]
    fn test_arrival_predictor() {
        let predictor = ArrivalPredictor::new();

        let p_time = predictor.predict_p_arrival(60.0, 10.0);
        let s_time = predictor.predict_s_arrival(60.0, 10.0);

        // S should arrive after P
        assert!(s_time > p_time);

        // S-P time increases with distance
        let sp_near = predictor.predict_sp_time(30.0, 10.0);
        let sp_far = predictor.predict_sp_time(100.0, 10.0);
        assert!(sp_far > sp_near);
    }

    #[test]
    fn test_arrival_predictor_from_model() {
        let model = sample_data::california_velocity_model();
        let predictor = ArrivalPredictor::from_velocity_model(&model);

        // Should use model velocities
        assert!(predictor.avg_vp > 4.0 && predictor.avg_vp < 8.0);
        assert!(predictor.avg_vs > 2.0 && predictor.avg_vs < 5.0);
    }

    #[test]
    fn test_distance_estimation() {
        let predictor = ArrivalPredictor::new();

        // Round-trip test
        let true_distance = 80.0;
        let sp_time = predictor.predict_sp_time(true_distance, 10.0);
        let estimated = predictor.estimate_distance(sp_time);

        // Should be approximately correct (ignoring depth effects)
        assert!((estimated - true_distance).abs() < 20.0);
    }

    #[test]
    fn test_ground_motion_prediction() {
        let config = sample_data::fast_config();
        let seismic_ai = SeismicAI::new(config);

        let source = sample_data::local_earthquake();
        let stations = sample_data::sparse_network();
        let velocity_model = sample_data::generic_velocity_model();

        let gm = seismic_ai.predict_ground_motion(&source, &stations, &velocity_model);

        assert_eq!(gm.len(), stations.len());

        for (station, motion) in &gm {
            assert!(motion.pga > 0.0);
            assert!(motion.p_arrival_time > 0.0);
            assert!(motion.s_arrival_time > motion.p_arrival_time);
            assert!(motion.mmi >= 1.0);
            // Station should be returned
            assert!(!station.code.is_empty());
        }
    }

    #[test]
    fn test_warning_generation() {
        let config = sample_data::fast_config();
        let mut seismic_ai = SeismicAI::new(config).with_warning_config(WarningConfig {
            min_magnitude: 3.0,
            min_stations: 2,
            ..Default::default()
        });

        let source = sample_data::local_earthquake();
        let stations = sample_data::sparse_network();
        let velocity_model = sample_data::generic_velocity_model();

        let warning = seismic_ai.generate_warning(&source, &stations, &velocity_model);

        // Should generate a warning for this magnitude
        assert!(!warning.event_id.is_empty());
        assert!(warning.detecting_stations >= 2);
        assert!(warning.confidence > 0.0);
    }

    #[test]
    fn test_no_warning_for_small_event() {
        let config = sample_data::fast_config();
        let mut seismic_ai = SeismicAI::new(config).with_warning_config(WarningConfig {
            min_magnitude: 5.0, // Higher threshold
            ..Default::default()
        });

        let mut source = sample_data::local_earthquake();
        source.magnitude = 3.0; // Below threshold

        let stations = sample_data::sparse_network();
        let velocity_model = sample_data::generic_velocity_model();

        let warning = seismic_ai.generate_warning(&source, &stations, &velocity_model);

        // Should not generate alert
        assert_eq!(warning.alert_level, AlertLevel::None);
    }

    #[test]
    fn test_simulation() {
        let (source, stations, velocity_model, config) = sample_data::local_test_scenario();
        let mut seismic_ai = SeismicAI::new(config);

        let result = seismic_ai.simulate(&source, &stations, &velocity_model);

        // Check result completeness
        assert_eq!(result.ground_motions.len(), stations.len());
        assert!(result.seismograms.is_some());
        assert!(result.snapshots.is_some());
        assert!(result.warning.is_some());
        assert!(result.computation_time_ms > 0.0);
    }

    #[test]
    fn test_seismogram_output() {
        let (source, stations, velocity_model, mut config) = sample_data::local_test_scenario();
        config.duration = 10.0; // Shorter for testing
        let mut seismic_ai = SeismicAI::new(config);

        let result = seismic_ai.simulate(&source, &stations, &velocity_model);

        let seismograms = result.seismograms.unwrap();
        assert_eq!(seismograms.len(), stations.len());

        for seis in &seismograms {
            assert!(!seis.station_code.is_empty());
            assert!(!seis.time.is_empty());
            assert_eq!(seis.east.len(), seis.time.len());
        }
    }

    #[test]
    fn test_wave_field_snapshot() {
        let (source, stations, velocity_model, mut config) = sample_data::local_test_scenario();
        config.duration = 5.0;
        let mut seismic_ai = SeismicAI::new(config);

        let result = seismic_ai.simulate(&source, &stations, &velocity_model);

        let snapshots = result.snapshots.unwrap();
        assert!(!snapshots.is_empty());

        for snap in &snapshots {
            assert!(!snap.displacement_x.is_empty());
            assert!(!snap.grid_x.is_empty());
            assert!(snap.time >= 0.0);
        }
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        assert!(result.computation_time_ms > 0.0);
        assert!(!result.ground_motions.is_empty());
    }

    #[test]
    fn test_attenuation_model_config() {
        let config = sample_data::fast_config();
        let seismic_ai = SeismicAI::new(config).with_attenuation_model(AttenuationModel::japan());

        assert_eq!(seismic_ai.attenuation_model.name(), "Japan");
    }

    #[test]
    fn test_warning_config() {
        let config = sample_data::fast_config();
        let warning_config = WarningConfig {
            min_magnitude: 4.0,
            min_stations: 5,
            ..Default::default()
        };
        let seismic_ai = SeismicAI::new(config).with_warning_config(warning_config);

        assert_eq!(seismic_ai.warning_config.min_magnitude, 4.0);
        assert_eq!(seismic_ai.warning_config.min_stations, 5);
    }

    #[test]
    fn test_major_scenario() {
        // Just test it runs without panicking
        let (source, stations, velocity_model, _) = sample_data::hayward_scenario();
        assert!(source.magnitude >= 6.5);
        assert!(stations.len() >= 50);
        assert!(!velocity_model.layers.is_empty());
    }
}