//! Composite risk-adjusted reward for PPO financial trading
//!
//! Combines multiple reward signals to balance return, risk, and drawdown:
//! R = w1 * `return_component` - w2 * `downside_risk` + w3 * `differential_return`
//!
//! Research: Risk-Aware RL Reward (2025) — 42% return, 8% max drawdown.

use serde::{Deserialize, Serialize};
use std::collections::VecDeque;

/// Configuration for composite reward
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompositeRewardConfig {
    /// Weight for return component (default: 1.0)
    pub return_weight: f64,
    /// Weight for downside deviation penalty (default: 0.5)
    pub downside_weight: f64,
    /// Weight for differential return (default: 0.3)
    pub differential_weight: f64,
    /// Window size for rolling computations (default: 20 steps)
    pub window_size: usize,
    /// Baseline SMA period for differential return (default: 50)
    pub baseline_period: usize,
}

impl Default for CompositeRewardConfig {
    fn default() -> Self {
        Self {
            return_weight: 1.0,
            downside_weight: 0.5,
            differential_weight: 0.3,
            window_size: 20,
            baseline_period: 50,
        }
    }
}

/// Composite risk-adjusted reward computer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompositeReward {
    config: CompositeRewardConfig,
    /// Rolling window of recent returns
    return_window: VecDeque<f64>,
    /// Longer window for baseline computation
    baseline_window: VecDeque<f64>,
    /// Peak portfolio value for drawdown tracking
    peak_value: f64,
    /// Current portfolio value
    current_value: f64,
}

impl CompositeReward {
    /// Create with default config
    #[must_use]
    pub fn new() -> Self {
        Self::with_config(CompositeRewardConfig::default())
    }

    /// Create with custom config
    #[must_use]
    pub fn with_config(config: CompositeRewardConfig) -> Self {
        Self {
            return_window: VecDeque::with_capacity(config.window_size),
            baseline_window: VecDeque::with_capacity(config.baseline_period),
            peak_value: 1.0, // Start at 1.0 (normalized)
            current_value: 1.0,
            config,
        }
    }

    /// Compute the composite reward for a single step
    ///
    /// # Arguments
    /// * `raw_return` - Single-step return (e.g., 0.001 for 0.1%)
    ///
    /// # Returns
    /// Composite reward combining all components
    pub fn compute(&mut self, raw_return: f32) -> f32 {
        let raw_return = f64::from(raw_return);
        // Update portfolio tracking
        self.current_value *= 1.0 + raw_return;
        if self.current_value > self.peak_value {
            self.peak_value = self.current_value;
        }

        // Update windows
        self.return_window.push_back(raw_return);
        if self.return_window.len() > self.config.window_size {
            self.return_window.pop_front();
        }
        self.baseline_window.push_back(raw_return);
        if self.baseline_window.len() > self.config.baseline_period {
            self.baseline_window.pop_front();
        }

        // Component 1: Return component (just the raw return, scaled)
        let return_component = raw_return * self.config.return_weight;

        // Component 2: Downside deviation penalty
        let downside = self.compute_downside_deviation();
        let downside_component = -downside * self.config.downside_weight;

        // Component 3: Differential return (vs baseline SMA)
        let differential = self.compute_differential_return(raw_return);
        let differential_component = differential * self.config.differential_weight;

        (return_component + downside_component + differential_component) as f32
    }

    /// Compute downside deviation (semi-variance of negative returns)
    fn compute_downside_deviation(&self) -> f64 {
        if self.return_window.len() < 2 {
            return 0.0;
        }
        let negative_returns: Vec<f64> = self.return_window.iter()
            .filter(|&&r| r < 0.0)
            .copied()
            .collect();
        if negative_returns.is_empty() {
            return 0.0;
        }
        let mean_neg: f64 = negative_returns.iter().sum::<f64>() / negative_returns.len() as f64;
        let variance: f64 = negative_returns
            .iter()
            .map(|r| (r - mean_neg).powi(2))
            .sum::<f64>()
            / negative_returns.len() as f64;
        variance.sqrt()
    }

    /// Compute differential return vs baseline SMA
    fn compute_differential_return(&self, current_return: f64) -> f64 {
        if self.baseline_window.len() < 5 {
            return 0.0;
        }
        let baseline_mean: f64 =
            self.baseline_window.iter().sum::<f64>() / self.baseline_window.len() as f64;
        current_return - baseline_mean
    }

    /// Get current drawdown from peak
    #[must_use]
    pub fn current_drawdown(&self) -> f64 {
        if self.peak_value <= 0.0 {
            return 0.0;
        }
        1.0 - (self.current_value / self.peak_value)
    }

    /// Reset state (call between episodes)
    pub fn reset(&mut self) {
        self.return_window.clear();
        self.baseline_window.clear();
        self.peak_value = 1.0;
        self.current_value = 1.0;
    }
}

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

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

    #[test]
    fn test_composite_reward_default() {
        let config = CompositeRewardConfig::default();
        assert!((config.return_weight - 1.0).abs() < f64::EPSILON);
        assert!((config.downside_weight - 0.5).abs() < f64::EPSILON);
        assert!((config.differential_weight - 0.3).abs() < f64::EPSILON);
        assert_eq!(config.window_size, 20);
        assert_eq!(config.baseline_period, 50);

        let reward = CompositeReward::new();
        assert!((reward.peak_value - 1.0).abs() < f64::EPSILON);
        assert!((reward.current_value - 1.0).abs() < f64::EPSILON);
        assert!(reward.return_window.is_empty());
        assert!(reward.baseline_window.is_empty());
    }

    #[test]
    fn test_composite_reward_positive_return() {
        let mut reward = CompositeReward::new();
        let r = reward.compute(0.01_f32); // 1% return
        // With only 1 data point, downside=0, differential=0 (baseline < 5)
        // So reward = 0.01 * 1.0 = 0.01
        assert!(r > 0.0, "Positive return should yield positive reward, got {r}");
        assert!(
            (r - 0.01_f32).abs() < 1e-6,
            "Single positive return should equal return_weight * raw_return"
        );
    }

    #[test]
    fn test_composite_reward_negative_return() {
        let mut reward = CompositeReward::new();
        let r = reward.compute(-0.02_f32); // -2% return
        // With only 1 data point, downside=0 (need >=2 in window), differential=0
        // So reward = -0.02 * 1.0 = -0.02
        assert!(r < 0.0, "Negative return should yield negative reward, got {r}");
    }

    #[test]
    fn test_composite_reward_downside_deviation() {
        let mut reward = CompositeReward::new();
        // Feed a sequence with several losses to build up downside deviation
        let returns: [f32; 10] = [
            -0.01, -0.02, -0.015, -0.005, 0.003, -0.01, -0.02, 0.001, -0.008, -0.012,
        ];
        let mut last_reward = 0.0_f32;
        for &ret in &returns {
            last_reward = reward.compute(ret);
        }
        // After many negative returns, downside deviation should be > 0
        // which means the penalty reduces the reward below the pure return component
        let pure_return_component = returns.last().copied().unwrap_or(0.0) * 1.0;
        assert!(
            last_reward < pure_return_component,
            "Downside deviation should penalize reward: got {last_reward} vs pure {pure_return_component}"
        );
    }

    #[test]
    fn test_composite_reward_differential() {
        let mut reward = CompositeReward::new();
        // Build a baseline of small returns
        for _ in 0..10 {
            reward.compute(0.001_f32); // 0.1% baseline
        }
        // Now a large positive return should get differential bonus
        let big_return = 0.05_f32; // 5% return — well above 0.1% baseline
        let r = reward.compute(big_return);
        // The differential component should be positive (outperforming baseline)
        let pure_return = big_return * 1.0;
        // r should be > pure_return because differential adds a bonus
        // (downside deviation may also contribute since we have no negative returns => 0 penalty)
        assert!(
            r > pure_return * 0.9,
            "Outperforming baseline should boost reward: got {r} vs pure {pure_return}"
        );
    }

    #[test]
    fn test_composite_reward_drawdown_tracking() {
        let mut reward = CompositeReward::new();
        // Go up then down
        reward.compute(0.10_f32); // Portfolio: 1.0 * 1.10 = 1.10
        assert!(
            (reward.current_drawdown()).abs() < 1e-6,
            "No drawdown at peak"
        );

        reward.compute(-0.05_f32); // Portfolio: 1.10 * 0.95 = 1.045
        let dd = reward.current_drawdown();
        assert!(dd > 0.0, "Should have drawdown after loss from peak");
        // Expected: 1 - 1.045/1.10 = 0.05 (tolerance accounts for f32 input promotion)
        assert!(
            (dd - 0.05).abs() < 1e-6,
            "Drawdown should be ~5%, got {dd}"
        );

        // Go to new peak
        reward.compute(0.10_f32); // Portfolio: 1.045 * 1.10 = 1.1495
        assert!(
            reward.current_drawdown().abs() < 1e-6,
            "Should be at new peak, no drawdown"
        );
    }

    #[test]
    fn test_composite_reward_reset() {
        let mut reward = CompositeReward::new();
        // Accumulate some state
        for i in 0..30 {
            reward.compute(if i % 2 == 0 { 0.01_f32 } else { -0.005_f32 });
        }
        assert!(!reward.return_window.is_empty());
        assert!(!reward.baseline_window.is_empty());

        reward.reset();
        assert!(reward.return_window.is_empty());
        assert!(reward.baseline_window.is_empty());
        assert!((reward.peak_value - 1.0).abs() < f64::EPSILON);
        assert!((reward.current_value - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_composite_reward_custom_config() {
        let config = CompositeRewardConfig {
            return_weight: 2.0,
            downside_weight: 1.0,
            differential_weight: 0.5,
            window_size: 10,
            baseline_period: 30,
        };
        let mut reward = CompositeReward::with_config(config);
        let r = reward.compute(0.01_f32);
        // With custom return_weight=2.0, single step: 0.01 * 2.0 = 0.02
        assert!(
            (r - 0.02_f32).abs() < 1e-6,
            "Custom return_weight should scale return: got {r}"
        );
    }

    #[test]
    fn test_composite_reward_all_positive() {
        let mut reward = CompositeReward::new();
        // All positive returns — no downside deviation penalty
        let mut last = 0.0_f32;
        for _ in 0..25 {
            last = reward.compute(0.005_f32);
        }
        // Downside deviation should be 0 since no negative returns
        // The differential component may be close to 0 (all returns are similar)
        // So reward should be close to 0.005 * 1.0 = 0.005
        assert!(
            last > 0.0,
            "All positive returns should yield positive reward"
        );
        // With 25 data points, baseline mean ~ 0.005, so differential ~ 0
        // Reward ~ 0.005 * 1.0 + 0 + 0 = 0.005
        assert!(
            (last - 0.005_f32).abs() < 0.001,
            "All-positive sequence should have reward near raw return, got {last}"
        );
    }

    #[test]
    fn test_composite_reward_mixed_sequence() {
        let mut reward = CompositeReward::new();
        // Realistic trading sequence: mostly small gains, occasional larger losses
        let returns: [f32; 20] = [
            0.002, 0.001, -0.003, 0.004, 0.001, -0.002, 0.003, -0.005, 0.002, 0.001, -0.001,
            0.003, 0.002, -0.004, 0.001, 0.002, -0.001, 0.003, -0.002, 0.001,
        ];
        let mut rewards = Vec::new();
        for &ret in &returns {
            rewards.push(reward.compute(ret));
        }
        // Verify we got rewards for every step
        assert_eq!(rewards.len(), returns.len());
        // Overall positive returns should yield a mostly positive reward sequence
        let positive_count = rewards.iter().filter(|&&r| r > 0.0).count();
        assert!(
            positive_count > rewards.len() / 3,
            "Mixed but net-positive sequence should have mostly positive rewards: {positive_count}/{}",
            rewards.len()
        );
        // Drawdown should be small for this mild sequence
        let dd = reward.current_drawdown();
        assert!(
            dd < 0.02,
            "Mild mixed sequence should have small drawdown, got {dd}"
        );
    }
}
