//! Percentile-based advantage scaling (`DreamerV3`)
//!
//! Tracks running P5/P95 of returns with EMA decay.
//! More robust than standard-deviation-based normalization for heavy-tailed
//! financial return distributions (fat tails, occasional outliers).
//!
//! # Why percentiles over std?
//! Standard deviation-based normalization assumes Gaussian tails. Financial returns
//! exhibit heavy tails (kurtosis >> 3), meaning rare large returns inflate std and
//! suppress the scaling of typical returns. Percentile ranges are robust to outliers.
//!
//! # References
//! - Hafner et al., "Mastering Diverse Domains through World Models" (`DreamerV3`, 2023)

use serde::{Deserialize, Serialize};

/// Tracks running percentiles via EMA for robust advantage scaling.
///
/// Maintains exponentially-weighted estimates of the 5th and 95th percentiles,
/// using them to normalize advantage values. The range `P95 - P5` captures 90%
/// of the data mass, making it robust to heavy tails.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PercentileScaler {
    /// Running estimate of 5th percentile
    p5: f64,
    /// Running estimate of 95th percentile
    p95: f64,
    /// EMA decay rate (default 0.99)
    decay: f64,
    /// Whether initialized with at least one batch
    initialized: bool,
    /// Minimum denominator to prevent division by zero
    min_scale: f64,
}

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

impl PercentileScaler {
    /// Create a new percentile scaler with default parameters.
    ///
    /// - `decay`: 0.99 (slow adaptation, stable estimates)
    /// - `min_scale`: 1.0 (prevents division by zero for constant data)
    #[must_use]
    pub const fn new() -> Self {
        Self {
            p5: 0.0,
            p95: 0.0,
            decay: 0.99,
            initialized: false,
            min_scale: 1.0,
        }
    }

    /// Create a percentile scaler with custom decay rate.
    ///
    /// # Arguments
    /// - `decay`: EMA decay in `(0, 1)`. Higher = slower adaptation.
    ///   Typical: 0.99 (stable) to 0.95 (reactive).
    /// - `min_scale`: Minimum P95-P5 range to prevent divide-by-zero.
    #[must_use]
    pub const fn with_params(decay: f64, min_scale: f64) -> Self {
        Self {
            p5: 0.0,
            p95: 0.0,
            decay: decay.clamp(0.0, 1.0),
            initialized: false,
            min_scale: min_scale.max(1e-8),
        }
    }

    /// Update running percentiles from a batch of values.
    ///
    /// Accepts `&[f32]` at the API boundary; converts to f64 internally for
    /// sorting and percentile computation to preserve numerical precision.
    ///
    /// Sorts the batch, computes exact P5/P95, then applies EMA update:
    /// `p_new = decay * p_old + (1 - decay) * p_actual`
    ///
    /// On first call, directly sets `p5`/`p95` to the actual percentiles
    /// (no EMA blending with zero).
    pub fn update(&mut self, values: &[f32]) {
        if values.is_empty() {
            return;
        }

        let mut sorted: Vec<f64> = values.iter().map(|&v| f64::from(v)).collect();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let n = sorted.len();
        let actual_p5 = Self::percentile_from_sorted(&sorted, n, 0.05);
        let actual_p95 = Self::percentile_from_sorted(&sorted, n, 0.95);

        if self.initialized {
            self.p5 = self.decay * self.p5 + (1.0 - self.decay) * actual_p5;
            self.p95 = self.decay * self.p95 + (1.0 - self.decay) * actual_p95;
        } else {
            self.p5 = actual_p5;
            self.p95 = actual_p95;
            self.initialized = true;
        }
    }

    /// Compute a percentile from a sorted slice using linear interpolation.
    fn percentile_from_sorted(sorted: &[f64], n: usize, p: f64) -> f64 {
        if n == 0 {
            return 0.0;
        }
        if n == 1 {
            return sorted.first().copied().unwrap_or(0.0);
        }

        // Linear interpolation index
        let idx = p * (n - 1) as f64;
        let lo = idx.floor() as usize;
        let hi = (lo + 1).min(n - 1);
        let frac = idx - lo as f64;

        let lo_val = sorted.get(lo).copied().unwrap_or(0.0);
        let hi_val = sorted.get(hi).copied().unwrap_or(lo_val);

        lo_val + frac * (hi_val - lo_val)
    }

    /// Scale a value by the running percentile range.
    ///
    /// Accepts f32, computes in f64 internally, returns f32.
    /// Returns `value / max(P95 - P5, min_scale)`.
    /// If not yet initialized, returns the value unchanged.
    #[inline]
    pub fn scale(&self, value: f32) -> f32 {
        if !self.initialized {
            return value;
        }
        let range = (self.p95 - self.p5).max(self.min_scale);
        (f64::from(value) / range) as f32
    }

    /// Scale a batch of f32 values in-place.
    ///
    /// Computes the range in f64, then divides each element.
    pub fn scale_batch(&self, values: &mut [f32]) {
        if !self.initialized {
            return;
        }
        let range = (self.p95 - self.p5).max(self.min_scale);
        for v in values.iter_mut() {
            *v = (f64::from(*v) / range) as f32;
        }
    }

    /// Get the current percentile range (P95 - P5).
    ///
    /// Returns `None` if not yet initialized.
    #[must_use]
    pub fn range(&self) -> Option<f64> {
        self.initialized.then(|| (self.p95 - self.p5).max(self.min_scale))
    }

    /// Get the current P5 estimate.
    #[must_use]
    pub const fn p5(&self) -> f64 {
        self.p5
    }

    /// Get the current P95 estimate.
    #[must_use]
    pub const fn p95(&self) -> f64 {
        self.p95
    }

    /// Whether the scaler has been initialized with at least one batch.
    #[must_use]
    pub const fn is_initialized(&self) -> bool {
        self.initialized
    }
}

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

    const EPSILON: f32 = 1e-6;
    const EPSILON_F64: f64 = 1e-6;

    #[test]
    fn test_percentile_scaler_new() {
        let scaler = PercentileScaler::new();
        assert!(!scaler.is_initialized());
        assert!(scaler.range().is_none());
        assert!((scaler.p5() - 0.0).abs() < EPSILON_F64);
        assert!((scaler.p95() - 0.0).abs() < EPSILON_F64);
    }

    #[test]
    fn test_percentile_scaler_first_update() {
        let mut scaler = PercentileScaler::new();
        // 100 values from 0..100
        let values: Vec<f32> = (0..100).map(|i| i as f32).collect();
        scaler.update(&values);

        assert!(scaler.is_initialized());
        // P5 of [0..99] ≈ 4.95, P95 ≈ 94.05
        assert!(
            (scaler.p5() - 4.95).abs() < 0.5,
            "P5 should be near 4.95, got {}",
            scaler.p5()
        );
        assert!(
            (scaler.p95() - 94.05).abs() < 0.5,
            "P95 should be near 94.05, got {}",
            scaler.p95()
        );
    }

    #[test]
    fn test_percentile_scaler_normal() {
        // Approximate normal distribution using sorted quantiles
        // For N(0,1): P5 ≈ -1.645, P95 ≈ +1.645
        let mut scaler = PercentileScaler::new();
        let n = 10000;
        // Use inverse CDF approximation: generate uniform then map
        let mut values = Vec::with_capacity(n);
        for i in 0..n {
            // Simple approximation: linearly spaced + some spread
            let u = (i as f64 + 0.5) / n as f64;
            // Approximate inverse normal CDF (Beasley-Springer-Moro)
            let z = approximate_inv_normal(u);
            values.push(z as f32);
        }
        scaler.update(&values);

        // P5 should be near -1.645, P95 near +1.645
        assert!(
            (scaler.p5() - (-1.645)).abs() < 0.1,
            "P5 should be near -1.645, got {}",
            scaler.p5()
        );
        assert!(
            (scaler.p95() - 1.645).abs() < 0.1,
            "P95 should be near 1.645, got {}",
            scaler.p95()
        );
    }

    #[test]
    fn test_percentile_scaler_heavy_tail() {
        // Heavy-tailed data: most values small, a few large outliers
        let mut scaler = PercentileScaler::new();
        let mut values: Vec<f32> = (0..990).map(|i| ((i as f32 - 495.0) / 495.0)).collect();
        // Add extreme outliers (1% each tail)
        for _ in 0..5 {
            values.push(-100.0_f32);
            values.push(100.0_f32);
        }
        scaler.update(&values);

        // P5/P95 should be set by the bulk, not the outliers
        let range = scaler.range().unwrap_or(0.0);
        // Range should be dominated by the [-1, 1] bulk, not the ±100 outliers
        assert!(
            range < 50.0,
            "percentile range should be robust to outliers, got {range}"
        );
    }

    #[test]
    fn test_percentile_scaler_min_scale() {
        // Constant values should not cause divide-by-zero
        let mut scaler = PercentileScaler::new();
        let values = vec![5.0_f32; 100];
        scaler.update(&values);

        // P5 == P95, so range falls back to min_scale
        let range = scaler.range().unwrap_or(0.0);
        assert!(
            range >= 1.0,
            "range should be at least min_scale=1.0, got {range}"
        );

        // Scaling should not produce Inf or NaN
        let scaled = scaler.scale(10.0_f32);
        assert!(scaled.is_finite(), "scaled value should be finite");
        assert!((scaled - 10.0_f32).abs() < EPSILON, "10.0 / 1.0 = 10.0, got {scaled}");
    }

    #[test]
    fn test_percentile_scaler_update_decay() {
        // Verify EMA decay behavior: second update blends with first
        let mut scaler = PercentileScaler::with_params(0.9, 1.0);

        // First batch: values in [0, 10]
        let batch1: Vec<f32> = (0..100).map(|i| i as f32 / 10.0).collect();
        scaler.update(&batch1);
        let p5_after_first = scaler.p5();
        let p95_after_first = scaler.p95();

        // Second batch: values in [100, 200]
        let batch2: Vec<f32> = (0..100).map(|i| 100.0_f32 + i as f32).collect();
        scaler.update(&batch2);

        // P5 should have moved toward batch2's P5 (≈ 104.95), but EMA dampens
        assert!(
            scaler.p5() > p5_after_first,
            "P5 should increase after batch2, was {p5_after_first}, now {}",
            scaler.p5()
        );
        // But not all the way to 104.95 (EMA with decay=0.9 keeps 90% of old)
        assert!(
            scaler.p5() < 104.95,
            "P5 should be EMA-damped, not jump to {}, got {}",
            104.95,
            scaler.p5()
        );

        // P95 should also have moved upward
        assert!(
            scaler.p95() > p95_after_first,
            "P95 should increase after batch2"
        );
    }

    #[test]
    fn test_percentile_scaler_financial_returns() {
        // Realistic trading return magnitudes:
        // Most returns are tiny (-0.001 to +0.001), occasional larger moves
        let mut scaler = PercentileScaler::new();
        let mut values: Vec<f32> = Vec::new();

        // 90% tiny returns in [-0.002, 0.002]
        for i in 0..900 {
            let r = (i as f32 - 450.0) / 450.0 * 0.002;
            values.push(r);
        }
        // 8% moderate returns in [-0.02, 0.02]
        for i in 0..80 {
            let r = (i as f32 - 40.0) / 40.0 * 0.02;
            values.push(r);
        }
        // 2% large returns in [-0.1, 0.1]
        for i in 0..20 {
            let r = (i as f32 - 10.0) / 10.0 * 0.1;
            values.push(r);
        }

        scaler.update(&values);

        // The raw P95-P5 spread is tiny for trading returns (~0.004)
        // but range() returns max(P95-P5, min_scale) so it floors at 1.0
        let raw_spread = scaler.p95() - scaler.p5();
        assert!(
            raw_spread < 0.2,
            "raw P95-P5 spread should be modest for trading returns, got {raw_spread}"
        );
        assert!(
            raw_spread > 0.001,
            "raw P95-P5 spread should be non-trivial, got {raw_spread}"
        );

        // range() applies min_scale floor
        let range = scaler.range().unwrap_or(0.0);
        assert!(
            range >= 1.0,
            "range should be at least min_scale, got {range}"
        );

        // Scaling a normal return should produce a reasonable magnitude
        let typical = scaler.scale(0.001_f32);
        assert!(
            typical.is_finite(),
            "scaled typical return should be finite"
        );
    }

    #[test]
    fn test_percentile_scaler_scale_batch() {
        let mut scaler = PercentileScaler::new();
        let values: Vec<f32> = (0..100).map(|i| i as f32).collect();
        scaler.update(&values);

        let mut batch = vec![10.0_f32, 20.0, 50.0];
        let range = scaler.range().unwrap_or(1.0);
        scaler.scale_batch(&mut batch);

        assert!(
            (f64::from(batch[0]) - 10.0 / range).abs() < EPSILON_F64,
            "batch[0] should be 10/range"
        );
        assert!(
            (f64::from(batch[1]) - 20.0 / range).abs() < EPSILON_F64,
            "batch[1] should be 20/range"
        );
        assert!(
            (f64::from(batch[2]) - 50.0 / range).abs() < EPSILON_F64,
            "batch[2] should be 50/range"
        );
    }

    #[test]
    fn test_percentile_scaler_empty_update() {
        let mut scaler = PercentileScaler::new();
        scaler.update(&[]);
        assert!(!scaler.is_initialized(), "empty update should not initialize");
    }

    #[test]
    fn test_percentile_scaler_single_value() {
        let mut scaler = PercentileScaler::new();
        scaler.update(&[42.0_f32]);
        assert!(scaler.is_initialized());
        // P5 == P95 == 42.0, range falls back to min_scale
        assert!((scaler.p5() - 42.0).abs() < EPSILON_F64);
        assert!((scaler.p95() - 42.0).abs() < EPSILON_F64);
    }

    #[test]
    fn test_percentile_scaler_uninitialized_passthrough() {
        let scaler = PercentileScaler::new();
        // Before initialization, scale should be passthrough
        assert!((scaler.scale(42.0_f32) - 42.0_f32).abs() < EPSILON);

        let mut batch = vec![1.0_f32, 2.0, 3.0];
        scaler.scale_batch(&mut batch);
        assert!((batch[0] - 1.0_f32).abs() < EPSILON);
        assert!((batch[1] - 2.0_f32).abs() < EPSILON);
        assert!((batch[2] - 3.0_f32).abs() < EPSILON);
    }

    // ---- Helper for approximate normal distribution ----

    /// Rational approximation of the inverse normal CDF (Abramowitz & Stegun 26.2.23)
    fn approximate_inv_normal(p: f64) -> f64 {
        if p <= 0.0 {
            return -6.0;
        }
        if p >= 1.0 {
            return 6.0;
        }

        let p_clamped = if p > 0.5 { 1.0 - p } else { p };
        let t = (-2.0 * p_clamped.ln()).sqrt();

        // Coefficients for the rational approximation
        let c0 = 2.515_517;
        let c1 = 0.802_853;
        let c2 = 0.010_328;
        let d1 = 1.432_788;
        let d2 = 0.189_269;
        let d3 = 0.001_308;

        let z = t - (c0 + c1 * t + c2 * t * t) / (1.0 + d1 * t + d2 * t * t + d3 * t * t * t);

        if p > 0.5 { z } else { -z }
    }
}
