//! Online reward normalization for PPO training
//!
//! Provides numerical stability for PPO reward calculations using
//! Exponential Moving Average (EMA) to prevent mean drift (BUG #41 fix).

use serde::{Deserialize, Serialize};

/// Online reward normalization using Exponential Moving Average (EMA)
///
/// Implements incremental mean and variance calculation to normalize rewards
/// to a standard normal distribution ~N(0,1). Uses EMA to prevent mean drift
/// in non-stationary markets (BUG #41 fix).
///
/// # Algorithm: Exponential Moving Average (EMA)
/// - Adapts to non-stationary reward distributions
/// - O(1) memory (no need to store all values)
/// - Exponentially decaying weights for old samples
///
/// # Why EMA over Welford's?
/// Welford's algorithm gives equal weight to all historical samples, causing
/// mean drift in non-stationary markets. When the market produces losses,
/// the mean drifts negative and stays there, inverting reward signs.
/// EMA gives exponentially decaying weight to old samples, adapting to regime changes.
///
/// # Usage
/// ```ignore
/// let mut normalizer = RewardNormalizer::new();
/// for reward in rewards {
///     normalizer.update(reward);
///     let normalized = normalizer.normalize(reward);
/// }
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RewardNormalizer {
    /// Running mean
    mean: f64,
    /// Running variance
    variance: f64,
    /// Mean decay rate (default: 0.01 = ~100-step effective window)
    alpha: f64,
    /// Variance decay rate (default: 0.01 = ~100-step effective window)
    beta: f64,
    /// Whether the normalizer has been initialized with at least one sample
    initialized: bool,
    /// Small constant for numerical stability
    epsilon: f64,
}

impl RewardNormalizer {
    /// Create a new reward normalizer with default decay rates
    ///
    /// # Default Parameters
    /// - alpha (mean decay): 0.01 = ~100-step effective window
    /// - beta (variance decay): 0.01 = ~100-step effective window
    ///
    /// # Effective Window Calculation
    /// For decay rate α, the effective window is approximately 1/α samples.
    /// Examples:
    /// - α=0.01 → 100 steps (balanced: reactive yet stable)
    /// - α=0.05 → 20 steps (aggressive: very reactive to changes)
    /// - α=0.001 → 1000 steps (conservative: smooth but slow)
    pub const fn new() -> Self {
        Self {
            mean: 0.0,
            variance: 1.0,
            alpha: 0.01,
            beta: 0.01,
            initialized: false,
            epsilon: 1e-8,
        }
    }

    /// Update running statistics with a new value (EMA algorithm)
    ///
    /// # Algorithm
    /// ```text
    /// First sample:
    ///   mean = value
    ///   variance = 1.0
    ///   initialized = true
    ///
    /// Subsequent samples:
    ///   mean = alpha * value + (1 - alpha) * mean
    ///   diff = value - mean
    ///   variance = beta * diff^2 + (1 - beta) * variance
    /// ```
    ///
    /// # Why This Prevents Mean Drift
    /// - Old samples decay exponentially: weight = (1-α)^t
    /// - After 100 steps (α=0.01): weight = 0.366 (63.4% decay)
    /// - After 200 steps: weight = 0.134 (86.6% decay)
    /// - After 500 steps: weight = 0.007 (99.3% decay)
    ///
    /// This allows the mean to "forget" old regimes and adapt to new ones.
    pub fn update(&mut self, value: f32) {
        let value = f64::from(value);
        if !self.initialized {
            // First sample: initialize mean and variance
            self.mean = value;
            self.variance = 1.0;
            self.initialized = true;
        } else {
            // EMA update for mean
            self.mean = self.alpha * value + (1.0 - self.alpha) * self.mean;

            // EMA update for variance
            let diff = value - self.mean;
            self.variance = self.beta * diff.powi(2) + (1.0 - self.beta) * self.variance;
        }
    }

    /// Normalize a value to standard normal distribution
    ///
    /// # Returns
    /// `(value - mean) / std` if initialized and std > epsilon, else value unchanged
    pub fn normalize(&self, value: f32) -> f32 {
        // Need at least one sample to have valid statistics
        if !self.initialized {
            return value;
        }

        let value = f64::from(value);
        let std = self.variance.sqrt();

        // Avoid division by zero for constant values
        if std < self.epsilon {
            return value as f32;
        }

        ((value - self.mean) / std) as f32
    }

    /// Get current mean and standard deviation
    pub fn get_stats(&self) -> (f64, f64) {
        if !self.initialized {
            return (0.0, 0.0);
        }

        let std = self.variance.sqrt();
        (self.mean, std)
    }

    /// Get count of values seen (backward compatibility)
    ///
    /// # Note
    /// EMA doesn't track sample count. Returns `u64::MAX` to indicate
    /// "many samples" for backward compatibility with code that calls this method.
    pub const fn count(&self) -> u64 {
        if self.initialized {
            u64::MAX
        } else {
            0
        }
    }
}

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