//! Adaptive Risk Management for DQN Training
//!
//! Helper methods for volatility-adjusted exploration, risk-adjusted rewards,
//! Kelly criterion position sizing, and risk tracker updates.

use tracing::debug;

use super::trainer::DQNTrainer;

impl DQNTrainer {
    // WAVE 16S: Adaptive Risk Management Helper Methods

    /// Calculate volatility-adjusted epsilon for exploration
    ///
    /// Adjusts epsilon based on recent return volatility:
    /// - Low volatility (<1%): reduce epsilon (exploit more)
    /// - High volatility (>5%): increase epsilon (explore more)
    /// - Moderate volatility: linear interpolation
    pub(crate) fn calculate_volatility_adjusted_epsilon(&self, base_epsilon: f64) -> f64 {
        if !self.hyperparams.enable_volatility_epsilon || self.volatility_returns.len() < 10 {
            return base_epsilon;
        }

        // Calculate volatility (standard deviation of returns)
        let mean: f64 = self.volatility_returns.iter().sum::<f64>() / self.volatility_returns.len() as f64;
        let variance: f64 = self.volatility_returns.iter()
            .map(|x| (x - mean).powi(2))
            .sum::<f64>() / self.volatility_returns.len() as f64;
        let volatility = variance.sqrt();

        // Adjust epsilon based on volatility regime
        let multiplier = if volatility < 0.01 {
            0.5  // Low volatility: exploit more (reduce epsilon)
        } else if volatility > 0.05 {
            2.0  // High volatility: explore more (increase epsilon)
        } else {
            // Linear interpolation between 0.01 and 0.05
            0.5 + (volatility - 0.01) / (0.05 - 0.01) * 1.5
        };

        let adjusted_epsilon = (base_epsilon * multiplier).clamp(0.01, 1.0);

        debug!(
            "Volatility-adjusted epsilon: base={:.4} × {:.2} = {:.4} (vol={:.4})",
            base_epsilon, multiplier, adjusted_epsilon, volatility
        );

        adjusted_epsilon
    }

    /// Calculate risk-adjusted reward using Sharpe-like normalization
    ///
    /// Divides raw reward by rolling volatility of recent trade returns to reward
    /// consistency over large but volatile returns. Uses trade_stats_history
    /// step_returns from GPU portfolio states.
    pub(crate) fn calculate_risk_adjusted_reward(&self, raw_reward: f64) -> f64 {
        if !self.hyperparams.enable_risk_adjusted_rewards {
            return raw_reward;
        }

        // Use step_returns from the most recent epoch's trade stats
        let recent_returns: Vec<f64> = self.trade_stats_history.iter()
            .rev()
            .take(3)
            .flat_map(|ts| ts.step_returns.iter().copied())
            .collect();

        if recent_returns.len() < 10 {
            return raw_reward;
        }

        // Calculate PnL volatility from recent step returns
        let n = recent_returns.len() as f64;
        let mean: f64 = recent_returns.iter().sum::<f64>() / n;
        let variance: f64 = recent_returns.iter()
            .map(|x| (x - mean).powi(2))
            .sum::<f64>() / n;
        let pnl_vol = variance.sqrt().max(1e-6);  // Prevent division by zero

        // Sharpe-like: reward / volatility
        raw_reward / pnl_vol
    }

    /// Get Kelly fraction for position sizing
    ///
    /// Returns Kelly criterion position size (0.0-0.25) based on trade history.
    /// Requires minimum trade history for statistical significance.
    pub fn get_kelly_fraction(&self) -> f64 {
        if !self.hyperparams.enable_kelly_sizing {
            return 1.0; // Full position sizing (disabled)
        }

        let kelly_opt = match &self.kelly_optimizer {
            Some(opt) => opt,
            None => return 1.0,
        };

        // Need minimum trades for Kelly calculation
        if self.trade_history.len() < self.hyperparams.kelly_min_trades {
            return 0.1; // Conservative default until enough history
        }

        // Calculate win/loss statistics
        let wins: Vec<f64> = self.trade_history.iter()
            .filter(|&&r| r > 0.0)
            .copied()
            .collect();
        let losses: Vec<f64> = self.trade_history.iter()
            .filter(|&&r| r < 0.0)
            .map(|r| -r)
            .collect();

        let win_prob = wins.len() as f64 / self.trade_history.len() as f64;
        let avg_win = if wins.is_empty() {
            0.01
        } else {
            wins.iter().sum::<f64>() / wins.len() as f64
        };
        let avg_loss = if losses.is_empty() {
            0.01
        } else {
            losses.iter().sum::<f64>() / losses.len() as f64
        };

        // Calculate Kelly fraction
        let kelly_result = kelly_opt.calculate_basic_kelly(win_prob, avg_win, avg_loss);
        let kelly_fraction = kelly_result.unwrap_or(0.1);

        // Apply fractional Kelly (conservative)
        let fractional_kelly = kelly_fraction * self.hyperparams.kelly_fractional;

        // Clamp to safety bounds
        fractional_kelly.clamp(0.01, self.hyperparams.kelly_max_fraction)
    }

    /// Update adaptive risk trackers with new market data
    pub(crate) fn update_risk_trackers(&mut self, reward: f64, price_return: f64) {
        // Update volatility tracker
        self.volatility_returns.push_back(price_return);
        if self.volatility_returns.len() > self.hyperparams.volatility_window {
            self.volatility_returns.pop_front();
        }

        // Update trade history for Kelly
        if reward.abs() > 1e-6 {  // Only track non-zero rewards
            self.trade_history.push_back(reward);
            if self.trade_history.len() > 500 {
                self.trade_history.pop_front();
            }
        }
    }
}
