//! Feature noise injection for robustness testing.
//!
//! Tests whether models are genuinely learning signal vs overfitting
//! to exact feature values by perturbing inputs and measuring degradation.

use crate::MLError;
use crate::types::TimeSeriesData;

/// Noise injection methods for robustness testing.
#[derive(Debug)]
pub struct NoiseInjector;

impl NoiseInjector {
    /// Add Gaussian noise to features. `sigma_fraction` is relative to each feature's
    /// standard deviation (e.g. 0.10 = 10% of std dev).
    /// Prices and timestamps are unchanged.
    pub fn gaussian_noise(
        data: &TimeSeriesData,
        sigma_fraction: f64,
        seed: u64,
    ) -> Result<TimeSeriesData, MLError> {
        let mut rng = SimpleRng::new(seed);
        let dim = data.features.first().map(|f| f.len()).unwrap_or(0);

        // Compute per-feature std dev
        let n = data.features.len() as f64;
        let mut means = vec![0.0_f64; dim];
        let mut vars = vec![0.0_f64; dim];

        for row in &data.features {
            for (j, val) in row.iter().enumerate() {
                if let Some(m) = means.get_mut(j) {
                    *m += f64::from(*val);
                }
            }
        }
        for m in &mut means {
            *m /= n;
        }

        for row in &data.features {
            for (j, val) in row.iter().enumerate() {
                let mean = means.get(j).copied().unwrap_or(0.0);
                if let Some(v) = vars.get_mut(j) {
                    *v += (f64::from(*val) - mean).powi(2);
                }
            }
        }
        // Use sample std dev; if zero (constant feature), fall back to the
        // absolute mean so that noise is proportional to the feature magnitude.
        // If both are zero, use 1.0 as a last-resort scale.
        let stds: Vec<f64> = vars
            .iter()
            .enumerate()
            .map(|(j, v)| {
                let s = (v / (n - 1.0).max(1.0)).sqrt();
                if s > 1e-15 {
                    s
                } else {
                    let abs_mean = means.get(j).copied().unwrap_or(0.0).abs();
                    if abs_mean > 1e-15 { abs_mean } else { 1.0 }
                }
            })
            .collect();

        let noisy_features: Vec<Vec<f32>> = data
            .features
            .iter()
            .map(|row| {
                row.iter()
                    .enumerate()
                    .map(|(j, &val)| {
                        let std = stds.get(j).copied().unwrap_or(1.0);
                        let noise = rng.normal() * sigma_fraction * std;
                        (f64::from(val) + noise) as f32
                    })
                    .collect()
            })
            .collect();

        TimeSeriesData::new(data.timestamps.clone(), noisy_features, data.prices.clone())
    }

    /// Randomly zero out features with the given probability.
    /// Prices and timestamps are unchanged.
    pub fn feature_dropout(
        data: &TimeSeriesData,
        drop_rate: f64,
        seed: u64,
    ) -> Result<TimeSeriesData, MLError> {
        let mut rng = SimpleRng::new(seed);

        let dropped_features: Vec<Vec<f32>> = data
            .features
            .iter()
            .map(|row| {
                row.iter()
                    .map(|&val| if rng.uniform() < drop_rate { 0.0 } else { val })
                    .collect()
            })
            .collect();

        TimeSeriesData::new(
            data.timestamps.clone(),
            dropped_features,
            data.prices.clone(),
        )
    }
}

/// Compute robustness score from Sharpe ratios at increasing noise levels.
///
/// `sharpes[0]` = clean, `sharpes[1..]` = at noise levels \[5%, 10%, 20%\].
/// Score = area under curve normalized to \[0, 1\].
/// 1.0 = no degradation, 0.0 = immediate collapse.
pub fn compute_robustness_score(sharpes: &[f64]) -> f64 {
    if sharpes.is_empty() {
        return 0.0;
    }
    let baseline = sharpes.first().copied().unwrap_or(0.0);
    if baseline.abs() < 1e-10 {
        return 0.0;
    }

    let n = sharpes.len() as f64;
    let area: f64 = sharpes.iter().map(|s| s / baseline).sum::<f64>() / n;
    area.clamp(0.0, 1.0)
}

/// Noise configuration for the validation harness.
#[derive(Debug, Clone)]
pub struct NoiseConfig {
    /// Noise levels to test (as fraction of feature std dev).
    pub noise_levels: Vec<f64>,
    /// RNG seed for reproducibility.
    pub seed: u64,
}

impl Default for NoiseConfig {
    fn default() -> Self {
        Self {
            noise_levels: vec![0.05, 0.10, 0.20],
            seed: 42,
        }
    }
}

/// Minimal PRNG (xorshift64) to avoid external dependency.
struct SimpleRng {
    state: u64,
}

impl SimpleRng {
    fn new(seed: u64) -> Self {
        Self {
            state: seed.wrapping_add(1),
        }
    }

    fn next_u64(&mut self) -> u64 {
        self.state ^= self.state << 13;
        self.state ^= self.state >> 7;
        self.state ^= self.state << 17;
        self.state
    }

    fn uniform(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1_u64 << 53) as f64
    }

    /// Box-Muller transform for normal distribution.
    fn normal(&mut self) -> f64 {
        let u1 = self.uniform().max(1e-15);
        let u2 = self.uniform();
        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
    }
}

#[cfg(test)]
#[allow(clippy::unreachable)]
mod tests {
    use super::*;
    use chrono::{TimeZone, Utc};
    use crate::types::TimeSeriesData;

    fn make_data(n: usize) -> TimeSeriesData {
        let timestamps: Vec<_> = (0..n)
            .map(|i| {
                Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0)
                    .single()
                    .unwrap_or_else(Utc::now)
                    + chrono::Duration::days(i as i64)
            })
            .collect();
        let features: Vec<Vec<f32>> = (0..n).map(|_| vec![1.0, 2.0, 3.0]).collect();
        let mut prices = Vec::with_capacity(n);
        let mut p = 100.0;
        for i in 0..n {
            prices.push(p);
            p += 0.01 + 0.005 * (i as f64).sin();
        }
        TimeSeriesData::new(timestamps, features, prices)
            .unwrap_or_else(|e| unreachable!("make_data should not fail: {e}"))
    }

    #[test]
    fn test_gaussian_noise_changes_features() {
        let data = make_data(50);
        let noisy = NoiseInjector::gaussian_noise(&data, 0.1, 42);
        assert!(noisy.is_ok());
        let noisy = noisy.unwrap_or_else(|_| unreachable!());
        let empty: Vec<f32> = vec![];
        let orig_f = data.features.first().unwrap_or(&empty);
        let noisy_f = noisy.features.first().unwrap_or(&empty);
        let any_different = orig_f
            .iter()
            .zip(noisy_f.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(
            any_different,
            "Gaussian noise should modify at least one feature"
        );
        assert_eq!(data.prices, noisy.prices);
    }

    #[test]
    fn test_feature_dropout_zeros_some_features() {
        let data = make_data(50);
        let dropped = NoiseInjector::feature_dropout(&data, 0.5, 42);
        assert!(dropped.is_ok());
        let dropped = dropped.unwrap_or_else(|_| unreachable!());
        let total_zeros: usize = dropped
            .features
            .iter()
            .flat_map(|row| row.iter())
            .filter(|v| v.abs() < 1e-10)
            .count();
        assert!(total_zeros > 0, "Dropout should zero out some features");
    }

    #[test]
    fn test_robustness_score_perfect_is_one() {
        let sharpes = vec![1.5, 1.5, 1.5, 1.5];
        let score = compute_robustness_score(&sharpes);
        assert!(
            (score - 1.0).abs() < 0.01,
            "Identical Sharpes should give score ~1.0, got {score}"
        );
    }

    #[test]
    fn test_robustness_score_collapse_is_low() {
        let sharpes = vec![2.0, 0.0, 0.0, 0.0];
        let score = compute_robustness_score(&sharpes);
        assert!(score < 0.5, "Collapse should give low score, got {score}");
    }
}
