//! Ground motion prediction equations and attenuation models.
//!
//! Implements various Ground Motion Prediction Equations (GMPEs) for estimating
//! ground motion intensity based on earthquake magnitude, distance, and site conditions.

use seismic_shared::{EarthquakeSource, GroundMotion, SiteClass, StationConfig, VelocityModel};

// ============================================================================
// GMPE Trait
// ============================================================================

/// Ground Motion Prediction Equation trait.
pub trait GMPE {
    /// Predict ground motion at a site given earthquake source and site parameters.
    fn predict(
        &self,
        source: &EarthquakeSource,
        station: &StationConfig,
        velocity_model: &VelocityModel,
    ) -> GroundMotion;

    /// Predict PGA in g.
    fn predict_pga(&self, magnitude: f64, distance_km: f64, depth_km: f64, vs30: f64) -> f64;

    /// Predict PGV in cm/s.
    fn predict_pgv(&self, magnitude: f64, distance_km: f64, depth_km: f64, vs30: f64) -> f64;

    /// Predict spectral acceleration at period T (in g).
    fn predict_sa(
        &self,
        magnitude: f64,
        distance_km: f64,
        depth_km: f64,
        vs30: f64,
        period: f64,
    ) -> f64;

    /// Get standard deviation (aleatory uncertainty).
    fn sigma(&self) -> f64;

    /// Get model name.
    fn name(&self) -> &str;
}

// ============================================================================
// Attenuation Model (Generic GMPE)
// ============================================================================

/// Generic attenuation model based on NGA-West2 style equations.
#[derive(Debug, Clone)]
pub struct AttenuationModel {
    /// Model name.
    name: String,
    /// Magnitude scaling coefficients.
    mag_coeffs: MagnitudeCoefficients,
    /// Distance scaling coefficients.
    dist_coeffs: DistanceCoefficients,
    /// Site coefficients.
    site_coeffs: SiteCoefficients,
    /// Depth coefficients.
    depth_coeffs: DepthCoefficients,
    /// Standard deviation (sigma).
    sigma: f64,
}

/// Magnitude scaling coefficients.
#[derive(Debug, Clone)]
pub struct MagnitudeCoefficients {
    /// Reference magnitude.
    pub m_ref: f64,
    /// Linear magnitude term.
    pub c1: f64,
    /// Quadratic magnitude term.
    pub c2: f64,
    /// Magnitude hinge point.
    pub m_hinge: f64,
    /// Slope below hinge.
    pub c3_low: f64,
    /// Slope above hinge.
    pub c3_high: f64,
}

impl Default for MagnitudeCoefficients {
    fn default() -> Self {
        Self {
            m_ref: 5.0,
            c1: 0.8,
            c2: 0.1,
            m_hinge: 6.5,
            c3_low: 0.3,
            c3_high: 0.15,
        }
    }
}

/// Distance scaling coefficients.
#[derive(Debug, Clone)]
pub struct DistanceCoefficients {
    /// Geometric spreading coefficient.
    pub c4: f64,
    /// Anelastic attenuation coefficient.
    pub c5: f64,
    /// Near-source saturation distance.
    pub h: f64,
    /// Transition distance.
    pub r_ref: f64,
}

impl Default for DistanceCoefficients {
    fn default() -> Self {
        Self {
            c4: -1.5,
            c5: -0.002,
            h: 5.0,
            r_ref: 1.0,
        }
    }
}

/// Site effect coefficients.
#[derive(Debug, Clone)]
pub struct SiteCoefficients {
    /// Reference Vs30 (m/s).
    pub vs30_ref: f64,
    /// Linear site term.
    pub c6: f64,
    /// Nonlinear site term reference PGA.
    pub pga_ref: f64,
    /// Nonlinear site coefficient.
    pub c7: f64,
    /// Basin depth coefficient (Z1.0).
    pub c8: f64,
    /// Basin depth coefficient (Z2.5).
    pub c9: f64,
}

impl Default for SiteCoefficients {
    fn default() -> Self {
        Self {
            vs30_ref: 760.0,
            c6: -0.5,
            pga_ref: 0.1,
            c7: -0.3,
            c8: 0.1,
            c9: 0.05,
        }
    }
}

/// Depth coefficients.
#[derive(Debug, Clone)]
pub struct DepthCoefficients {
    /// Reference depth (km).
    pub z_ref: f64,
    /// Depth scaling.
    pub c10: f64,
    /// Maximum depth effect.
    pub z_max: f64,
}

impl Default for DepthCoefficients {
    fn default() -> Self {
        Self {
            z_ref: 10.0,
            c10: -0.02,
            z_max: 30.0,
        }
    }
}

impl AttenuationModel {
    /// Create a new attenuation model.
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            mag_coeffs: MagnitudeCoefficients::default(),
            dist_coeffs: DistanceCoefficients::default(),
            site_coeffs: SiteCoefficients::default(),
            depth_coeffs: DepthCoefficients::default(),
            sigma: 0.7,
        }
    }

    /// Create NGA-West2 style model.
    pub fn nga_west2() -> Self {
        Self::new("NGA-West2")
    }

    /// Create a California-specific model.
    pub fn california() -> Self {
        let mut model = Self::new("California");
        model.dist_coeffs.c5 = -0.0025; // Higher anelastic attenuation
        model.site_coeffs.c6 = -0.6; // Stronger site effects
        model
    }

    /// Create a Japan-style model.
    pub fn japan() -> Self {
        let mut model = Self::new("Japan");
        model.dist_coeffs.c5 = -0.003; // Higher attenuation in volcanic regions
        model.depth_coeffs.c10 = -0.03; // Stronger depth effects
        model
    }

    /// Calculate rupture distance.
    fn rupture_distance(&self, source: &EarthquakeSource, station: &StationConfig) -> f64 {
        let epicentral = source.hypocenter.location.distance_km(&station.location);
        let depth = source.hypocenter.depth_km;

        // Approximate rupture distance (Rrup)
        let rup_length = source.estimated_rupture_length();

        if epicentral < rup_length {
            // Close to rupture
            depth
        } else {
            // Far from rupture
            ((epicentral - rup_length / 2.0).powi(2) + depth.powi(2)).sqrt()
        }
    }

    /// Calculate Joyner-Boore distance.
    fn joyner_boore_distance(&self, source: &EarthquakeSource, station: &StationConfig) -> f64 {
        let epicentral = source.hypocenter.location.distance_km(&station.location);
        let rup_length = source.estimated_rupture_length();

        // Rjb = max(0, epicentral - rupture_half_length)
        (epicentral - rup_length / 2.0).max(0.0)
    }

    /// Calculate P-wave arrival time.
    fn p_arrival_time(
        &self,
        distance_km: f64,
        depth_km: f64,
        velocity_model: &VelocityModel,
    ) -> f64 {
        velocity_model.p_wave_travel_time(distance_km, depth_km)
    }

    /// Calculate S-wave arrival time.
    fn s_arrival_time(
        &self,
        distance_km: f64,
        depth_km: f64,
        velocity_model: &VelocityModel,
    ) -> f64 {
        velocity_model.s_wave_travel_time(distance_km, depth_km)
    }

    /// Calculate magnitude scaling term.
    fn magnitude_term(&self, magnitude: f64) -> f64 {
        let mc = &self.mag_coeffs;
        let dm = magnitude - mc.m_ref;

        if magnitude < mc.m_hinge {
            mc.c1 * dm + mc.c2 * dm.powi(2) + mc.c3_low * (magnitude - mc.m_hinge).max(0.0)
        } else {
            mc.c1 * dm + mc.c2 * dm.powi(2) + mc.c3_high * (magnitude - mc.m_hinge)
        }
    }

    /// Calculate distance scaling term.
    fn distance_term(&self, r_rup: f64, magnitude: f64) -> f64 {
        let dc = &self.dist_coeffs;

        // Magnitude-dependent saturation
        let h = dc.h + 0.5 * (magnitude - 5.0).max(0.0);
        let r = (r_rup.powi(2) + h.powi(2)).sqrt();

        // Geometric spreading + anelastic attenuation
        dc.c4 * (r / dc.r_ref).ln() + dc.c5 * r
    }

    /// Calculate site amplification term.
    fn site_term(&self, vs30: f64, pga_rock: f64) -> f64 {
        let sc = &self.site_coeffs;

        // Linear site response
        let linear = sc.c6 * (vs30 / sc.vs30_ref).ln();

        // Nonlinear site response (for soft sites at high PGA)
        let nonlinear = if vs30 < sc.vs30_ref && pga_rock > sc.pga_ref {
            sc.c7
                * ((pga_rock + sc.pga_ref) / (2.0 * sc.pga_ref)).ln()
                * ((vs30 / sc.vs30_ref).min(1.0))
        } else {
            0.0
        };

        linear + nonlinear
    }

    /// Calculate depth term.
    fn depth_term(&self, depth_km: f64) -> f64 {
        let dc = &self.depth_coeffs;
        let d = depth_km.min(dc.z_max);
        dc.c10 * (d - dc.z_ref)
    }

    /// Calculate basin depth term.
    fn basin_term(&self, station: &StationConfig) -> f64 {
        let sc = &self.site_coeffs;

        let z1_term = station.z1_0.map_or(0.0, |z| sc.c8 * z.ln());
        let z25_term = station.z2_5.map_or(0.0, |z| sc.c9 * z.ln());

        z1_term + z25_term
    }

    /// Period-dependent coefficients for spectral acceleration.
    fn period_coefficients(&self, period: f64) -> (f64, f64, f64) {
        // Simplified period-dependent scaling
        // (amplitude factor, magnitude term scale, distance term scale)
        if period < 0.1 {
            (1.0, 1.0, 1.0)
        } else if period < 0.5 {
            (1.0 + 0.5 * (period / 0.5).ln(), 1.1, 0.95)
        } else if period < 1.0 {
            (1.2 - 0.3 * (period - 0.5), 1.15, 0.9)
        } else if period < 3.0 {
            (0.9 - 0.2 * (period - 1.0), 1.2, 0.85)
        } else {
            (0.5, 1.25, 0.8)
        }
    }
}

impl GMPE for AttenuationModel {
    fn predict(
        &self,
        source: &EarthquakeSource,
        station: &StationConfig,
        velocity_model: &VelocityModel,
    ) -> GroundMotion {
        let r_rup = self.rupture_distance(source, station);
        let _r_jb = self.joyner_boore_distance(source, station);
        let epicentral = source.hypocenter.location.distance_km(&station.location);
        let vs30 = station.get_vs30();

        // Predict PGA and PGV
        let pga = self.predict_pga(source.magnitude, r_rup, source.hypocenter.depth_km, vs30);
        let pgv = self.predict_pgv(source.magnitude, r_rup, source.hypocenter.depth_km, vs30);

        // Calculate arrival times
        let p_arrival = self.p_arrival_time(epicentral, source.hypocenter.depth_km, velocity_model);
        let s_arrival = self.s_arrival_time(epicentral, source.hypocenter.depth_km, velocity_model);

        // Predict spectral accelerations
        let sa_03 = self.predict_sa(
            source.magnitude,
            r_rup,
            source.hypocenter.depth_km,
            vs30,
            0.3,
        );
        let sa_10 = self.predict_sa(
            source.magnitude,
            r_rup,
            source.hypocenter.depth_km,
            vs30,
            1.0,
        );
        let sa_30 = self.predict_sa(
            source.magnitude,
            r_rup,
            source.hypocenter.depth_km,
            vs30,
            3.0,
        );

        // Estimate duration
        let duration_5_95 =
            5.0 + 0.5 * (s_arrival - p_arrival) + 2.0 * (source.magnitude - 5.0).max(0.0);

        // Estimate PGD
        let pgd = pgv * 0.1 * (source.magnitude - 4.0).max(1.0);

        // Calculate MMI
        let mmi = f64::midpoint(
            GroundMotion::estimate_mmi_from_pga(pga),
            GroundMotion::estimate_mmi_from_pgv(pgv),
        );

        GroundMotion {
            pga,
            pgv,
            pgd,
            p_arrival_time: p_arrival,
            s_arrival_time: s_arrival,
            duration_5_95,
            sa_03,
            sa_10,
            sa_30,
            mmi,
        }
    }

    fn predict_pga(&self, magnitude: f64, distance_km: f64, depth_km: f64, vs30: f64) -> f64 {
        // Calculate log10(PGA) in g
        let mag_term = self.magnitude_term(magnitude);
        let dist_term = self.distance_term(distance_km, magnitude);
        let depth_term = self.depth_term(depth_km);

        // First pass: rock site PGA
        let ln_pga_rock = mag_term + dist_term + depth_term;
        let pga_rock = ln_pga_rock.exp();

        // Add site effects
        let site_term = self.site_term(vs30, pga_rock);

        let ln_pga = ln_pga_rock + site_term;

        // Convert from ln to actual value, with realistic scaling
        let pga = (ln_pga - 4.0).exp(); // Base adjustment

        // Clamp to reasonable range
        pga.clamp(0.0001, 2.0)
    }

    fn predict_pgv(&self, magnitude: f64, distance_km: f64, depth_km: f64, vs30: f64) -> f64 {
        // PGV scales differently than PGA
        let pga = self.predict_pga(magnitude, distance_km, depth_km, vs30);

        // Empirical PGV-PGA relationship
        // PGV (cm/s) ~ 100 * PGA (g) for typical earthquakes
        // with magnitude-dependent adjustment
        let mag_factor = 0.8 + 0.1 * (magnitude - 5.0).max(0.0);

        100.0 * pga * mag_factor
    }

    fn predict_sa(
        &self,
        magnitude: f64,
        distance_km: f64,
        depth_km: f64,
        vs30: f64,
        period: f64,
    ) -> f64 {
        let pga = self.predict_pga(magnitude, distance_km, depth_km, vs30);
        let (amp_factor, mag_scale, dist_scale) = self.period_coefficients(period);

        // Adjust for period
        let period_amp = if period < 0.1 {
            1.0
        } else if period < 1.0 {
            // Short periods amplified
            1.0 + 1.5 * (1.0 - (period - 0.1) / 0.9)
        } else {
            // Long periods attenuated
            1.0 / (1.0 + period - 1.0)
        };

        pga * amp_factor * period_amp * mag_scale * dist_scale
    }

    fn sigma(&self) -> f64 {
        self.sigma
    }

    fn name(&self) -> &str {
        &self.name
    }
}

// ============================================================================
// Site Effects
// ============================================================================

/// Site effects amplification model.
#[derive(Debug, Clone)]
pub struct SiteEffects {
    /// Site class.
    pub site_class: SiteClass,
    /// Measured Vs30 (m/s).
    pub vs30: f64,
    /// Basin depth Z1.0 (km).
    pub z1_0: Option<f64>,
    /// Basin depth Z2.5 (km).
    pub z2_5: Option<f64>,
    /// Predominant period (seconds).
    pub predominant_period: Option<f64>,
}

impl SiteEffects {
    /// Create site effects from station config.
    pub fn from_station(station: &StationConfig) -> Self {
        Self {
            site_class: station.site_class,
            vs30: station.get_vs30(),
            z1_0: station.z1_0,
            z2_5: station.z2_5,
            predominant_period: None,
        }
    }

    /// Calculate frequency-dependent amplification.
    pub fn amplification(&self, frequency: f64) -> f64 {
        let base_amp = self.site_class.amplification_factor();

        // Simple resonance model
        let resonance = if let Some(t0) = self.predominant_period {
            let f0 = 1.0 / t0;
            let ratio = frequency / f0;

            // Amplification peak at resonance
            if ratio > 0.5 && ratio < 2.0 {
                1.0 + 0.5 * (1.0 - (ratio - 1.0).abs())
            } else {
                1.0
            }
        } else {
            // Estimate from Vs30
            let f0 = self.vs30 / (4.0 * 30.0); // Quarter-wavelength approximation
            let ratio = frequency / f0;

            if ratio > 0.5 && ratio < 2.0 {
                1.0 + 0.3 * (1.0 - (ratio - 1.0).abs())
            } else {
                1.0
            }
        };

        base_amp * resonance
    }

    /// Calculate nonlinear site response for high input motion.
    pub fn nonlinear_factor(&self, pga_input: f64) -> f64 {
        if pga_input < 0.1 {
            // Linear regime
            1.0
        } else if pga_input < 0.3 {
            // Transition
            1.0 - 0.2 * (pga_input - 0.1) / 0.2
        } else {
            // Nonlinear (soft sites de-amplify at high shaking)
            0.8 - 0.3 * (pga_input - 0.3).min(0.5)
        }
    }

    /// Apply site effects to ground motion.
    pub fn apply(&self, input: &GroundMotion) -> GroundMotion {
        let amp = self.site_class.amplification_factor();
        let nonlinear = self.nonlinear_factor(input.pga);

        GroundMotion {
            pga: input.pga * amp * nonlinear,
            pgv: input.pgv * amp * nonlinear.sqrt(), // PGV less affected by nonlinearity
            pgd: input.pgd * amp,
            p_arrival_time: input.p_arrival_time,
            s_arrival_time: input.s_arrival_time,
            duration_5_95: input.duration_5_95 * (1.0 + 0.2 * (amp - 1.0)),
            sa_03: input.sa_03 * amp * self.amplification(1.0 / 0.3) * nonlinear,
            sa_10: input.sa_10 * amp * self.amplification(1.0) * nonlinear.sqrt(),
            sa_30: input.sa_30 * amp * self.amplification(1.0 / 3.0),
            mmi: GroundMotion::estimate_mmi_from_pga(input.pga * amp * nonlinear),
        }
    }
}

impl Default for SiteEffects {
    fn default() -> Self {
        Self {
            site_class: SiteClass::C,
            vs30: 500.0,
            z1_0: None,
            z2_5: None,
            predominant_period: None,
        }
    }
}

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

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

    #[test]
    fn test_attenuation_model_creation() {
        let model = AttenuationModel::nga_west2();
        assert_eq!(model.name(), "NGA-West2");

        let ca_model = AttenuationModel::california();
        assert_eq!(ca_model.name(), "California");

        let jp_model = AttenuationModel::japan();
        assert_eq!(jp_model.name(), "Japan");
    }

    #[test]
    fn test_pga_magnitude_scaling() {
        let model = AttenuationModel::nga_west2();

        let pga_m5 = model.predict_pga(5.0, 10.0, 10.0, 760.0);
        let pga_m6 = model.predict_pga(6.0, 10.0, 10.0, 760.0);
        let pga_m7 = model.predict_pga(7.0, 10.0, 10.0, 760.0);

        // PGA should increase with magnitude
        assert!(pga_m6 > pga_m5);
        assert!(pga_m7 > pga_m6);
    }

    #[test]
    fn test_pga_distance_scaling() {
        let model = AttenuationModel::nga_west2();

        let pga_10km = model.predict_pga(6.0, 10.0, 10.0, 760.0);
        let pga_50km = model.predict_pga(6.0, 50.0, 10.0, 760.0);
        let pga_100km = model.predict_pga(6.0, 100.0, 10.0, 760.0);

        // PGA should decrease with distance
        assert!(pga_50km < pga_10km);
        assert!(pga_100km < pga_50km);
    }

    #[test]
    fn test_pga_site_effects() {
        let model = AttenuationModel::nga_west2();

        let pga_rock = model.predict_pga(6.0, 20.0, 10.0, 760.0);
        let pga_soft = model.predict_pga(6.0, 20.0, 10.0, 200.0);

        // Soft site should have higher PGA (in linear regime)
        // Note: This depends on the model coefficients
        // Just check that they're different
        assert!((pga_rock - pga_soft).abs() > 0.0001);
    }

    #[test]
    fn test_pgv_prediction() {
        let model = AttenuationModel::nga_west2();

        let pgv = model.predict_pgv(6.0, 20.0, 10.0, 760.0);

        // PGV should be positive and reasonable
        assert!(pgv > 0.0);
        assert!(pgv < 200.0); // Reasonable upper bound for M6 at 20km
    }

    #[test]
    fn test_spectral_acceleration() {
        let model = AttenuationModel::nga_west2();

        let pga = model.predict_pga(6.0, 20.0, 10.0, 760.0);
        let sa_03 = model.predict_sa(6.0, 20.0, 10.0, 760.0, 0.3);
        let sa_10 = model.predict_sa(6.0, 20.0, 10.0, 760.0, 1.0);
        let sa_30 = model.predict_sa(6.0, 20.0, 10.0, 760.0, 3.0);

        // SA at 0.3s typically higher than PGA
        assert!(sa_03 > pga * 0.5);

        // Long period SA should be lower
        assert!(sa_30 < sa_03);
    }

    #[test]
    fn test_ground_motion_prediction() {
        let model = AttenuationModel::nga_west2();
        let source = seismic_shared::sample_local_earthquake();
        let station = seismic_shared::StationConfig::new("TEST", 37.9, -122.3);
        let velocity_model = seismic_shared::sample_california_velocity_model();

        let gm = model.predict(&source, &station, &velocity_model);

        // Check all values are positive and finite
        assert!(gm.pga > 0.0 && gm.pga.is_finite());
        assert!(gm.pgv > 0.0 && gm.pgv.is_finite());
        assert!(gm.p_arrival_time > 0.0);
        assert!(gm.s_arrival_time > gm.p_arrival_time);
        assert!(gm.mmi >= 1.0 && gm.mmi <= 12.0);
    }

    #[test]
    fn test_site_effects() {
        let site = SiteEffects {
            site_class: SiteClass::E,
            vs30: 150.0,
            ..Default::default()
        };

        // Soft soil amplification
        assert!(site.amplification(1.0) > 1.0);

        // Nonlinear de-amplification at high PGA
        let nlin_low = site.nonlinear_factor(0.05);
        let nlin_high = site.nonlinear_factor(0.4);
        assert!(nlin_high < nlin_low);
    }

    #[test]
    fn test_site_effects_application() {
        let site = SiteEffects {
            site_class: SiteClass::D,
            vs30: 270.0,
            ..Default::default()
        };

        let input = GroundMotion {
            pga: 0.1,
            pgv: 10.0,
            pgd: 1.0,
            p_arrival_time: 5.0,
            s_arrival_time: 10.0,
            duration_5_95: 15.0,
            sa_03: 0.2,
            sa_10: 0.15,
            sa_30: 0.05,
            mmi: 6.0,
        };

        let output = site.apply(&input);

        // Amplification for site class D
        assert!(output.pga > input.pga);
        assert!(output.pgv > input.pgv);
        // Arrival times unchanged
        assert_eq!(output.p_arrival_time, input.p_arrival_time);
    }

    #[test]
    fn test_sigma() {
        let model = AttenuationModel::nga_west2();
        let sigma = model.sigma();

        // Typical GMPE sigma is 0.5-0.9
        assert!(sigma > 0.4 && sigma < 1.0);
    }

    #[test]
    fn test_magnitude_coefficients() {
        let mc = MagnitudeCoefficients::default();
        assert!(mc.m_ref > 0.0);
        assert!(mc.m_hinge > mc.m_ref);
    }

    #[test]
    fn test_distance_coefficients() {
        let dc = DistanceCoefficients::default();
        assert!(dc.c4 < 0.0); // Geometric spreading is negative
        assert!(dc.c5 < 0.0); // Anelastic attenuation is negative
    }
}