//! PPO reward shaping for financial trading
//!
//! Provides additional reward components beyond raw `PnL` to guide PPO training:
//! - Hold penalty: discourages staying flat when signals exist
//! - Rolling Sharpe: risk-adjusted reward component
//! - Diversity bonus: encourages varied action selection (smooth quadratic)

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

/// Number of factored actions in PPO action space (5 exposure x 3 order x 3 urgency).
const NUM_FACTORED_ACTIONS: usize = 45;

/// Minimum return window size before Sharpe computation kicks in.
const MIN_SHARPE_WINDOW: usize = 5;

/// Minimum total actions before diversity bonus kicks in.
const MIN_DIVERSITY_ACTIONS: u32 = 10;

/// PPO reward shaper with configurable components.
///
/// Combines multiple reward signals to guide PPO training beyond raw `PnL`:
///
/// 1. **Hold penalty** — applied when the agent is flat and signals suggest
///    an actionable opportunity.  Discourages inactivity.
/// 2. **Rolling Sharpe** — risk-adjusted return component computed over a
///    sliding window of recent returns.
/// 3. **Diversity bonus** — smooth quadratic bonus based on Shannon entropy
///    of the action distribution.  Encourages exploration.
///
/// The shaped reward is:
/// ```text
/// shaped = raw_reward
///        - hold_penalty_weight   * I(flat & signal)
///        + sharpe_weight         * rolling_sharpe
///        + diversity_weight      * (entropy / max_entropy)^2
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PPORewardShaper {
    /// Hold penalty weight (applied when position is flat).
    hold_penalty_weight: f64,
    /// Rolling Sharpe weight.
    sharpe_weight: f64,
    /// Diversity bonus weight (smooth quadratic entropy bonus).
    diversity_weight: f64,
    /// Rolling window of recent returns for Sharpe computation.
    return_window: VecDeque<f64>,
    /// Window size for rolling Sharpe (default: 20).
    window_size: usize,
    /// Action frequency tracker for diversity (indexed by action id).
    action_counts: Vec<u32>,
    /// Total actions observed for diversity computation.
    total_actions: u32,
}

impl PPORewardShaper {
    /// Create a new reward shaper with configurable weights.
    ///
    /// # Arguments
    /// * `hold_penalty_weight` — penalty magnitude when flat with signal
    /// * `sharpe_weight` — scaling factor for rolling Sharpe component
    /// * `diversity_weight` — scaling factor for quadratic entropy bonus
    pub fn new(hold_penalty_weight: f64, sharpe_weight: f64, diversity_weight: f64) -> Self {
        Self {
            hold_penalty_weight,
            sharpe_weight,
            diversity_weight,
            return_window: VecDeque::with_capacity(20),
            window_size: 20,
            action_counts: vec![0; NUM_FACTORED_ACTIONS],
            total_actions: 0,
        }
    }

    /// Shape a reward with all enabled components.
    ///
    /// # Arguments
    /// * `raw_reward` — raw `PnL` reward from the environment
    /// * `is_flat` — whether the agent currently holds no position
    /// * `has_signal` — whether market signals suggest an actionable opportunity
    /// * `action_index` — index of the action taken (0..44)
    ///
    /// # Returns
    /// `shaped = raw_reward + hold_penalty + sharpe_component + diversity_bonus`
    pub fn shape_reward(
        &mut self,
        raw_reward: f32,
        is_flat: bool,
        has_signal: bool,
        action_index: usize,
    ) -> f32 {
        let raw_reward = f64::from(raw_reward);
        let mut shaped = raw_reward;

        // 1. Hold penalty: penalize flat position when signals exist
        if is_flat && has_signal && self.hold_penalty_weight > 0.0 {
            shaped -= self.hold_penalty_weight;
        }

        // 2. Rolling Sharpe component
        self.return_window.push_back(raw_reward);
        if self.return_window.len() > self.window_size {
            self.return_window.pop_front();
        }
        if self.return_window.len() >= MIN_SHARPE_WINDOW && self.sharpe_weight > 0.0 {
            let sharpe = self.compute_rolling_sharpe();
            shaped += self.sharpe_weight * sharpe;
        }

        // 3. Diversity bonus (smooth quadratic entropy)
        if let Some(count) = self.action_counts.get_mut(action_index) {
            *count += 1;
        }
        self.total_actions += 1;
        if self.diversity_weight > 0.0 && self.total_actions > MIN_DIVERSITY_ACTIONS {
            let entropy = self.compute_action_entropy();
            let max_entropy = (NUM_FACTORED_ACTIONS as f64).ln();
            let normalized = if max_entropy > 0.0 {
                (entropy / max_entropy).clamp(0.0, 1.0)
            } else {
                0.0
            };
            // Smooth quadratic bonus: peaks at max entropy, zero at min
            shaped += self.diversity_weight * normalized * normalized;
        }

        shaped as f32
    }

    /// Compute rolling Sharpe ratio from return window.
    fn compute_rolling_sharpe(&self) -> f64 {
        if self.return_window.is_empty() {
            return 0.0;
        }
        let n = self.return_window.len() as f64;
        let mean: f64 = self.return_window.iter().sum::<f64>() / n;
        let variance: f64 = self
            .return_window
            .iter()
            .map(|r| (r - mean).powi(2))
            .sum::<f64>()
            / n;
        let std = variance.sqrt();
        if std < 1e-10 {
            return 0.0;
        }
        mean / std
    }

    /// Compute Shannon entropy of action distribution.
    fn compute_action_entropy(&self) -> f64 {
        if self.total_actions == 0 {
            return 0.0;
        }
        let total = self.total_actions as f64;
        self.action_counts
            .iter()
            .filter(|&&c| c > 0)
            .map(|&c| {
                let p = c as f64 / total;
                -p * p.ln()
            })
            .sum()
    }

    /// Reset internal state (call between episodes).
    pub fn reset(&mut self) {
        self.return_window.clear();
        self.action_counts.fill(0);
        self.total_actions = 0;
    }

    /// Get current rolling Sharpe (for logging / diagnostics).
    pub fn current_sharpe(&self) -> f64 {
        self.compute_rolling_sharpe()
    }

    /// Get hold penalty weight.
    pub const fn hold_penalty_weight(&self) -> f64 {
        self.hold_penalty_weight
    }

    /// Get Sharpe weight.
    pub const fn sharpe_weight(&self) -> f64 {
        self.sharpe_weight
    }

    /// Get diversity weight.
    pub const fn diversity_weight(&self) -> f64 {
        self.diversity_weight
    }
}

impl Default for PPORewardShaper {
    fn default() -> Self {
        Self::new(0.01, 0.1, 0.05) // Modest defaults
    }
}

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

    #[test]
    fn test_reward_shaping_no_penalty_when_not_flat() {
        let mut shaper = PPORewardShaper::new(0.5, 0.0, 0.0);
        // Not flat => no hold penalty even with signal
        let shaped = shaper.shape_reward(1.0_f32, false, true, 0);
        assert!((shaped - 1.0_f32).abs() < 1e-6, "Expected 1.0, got {}", shaped);
    }

    #[test]
    fn test_reward_shaping_hold_penalty() {
        let mut shaper = PPORewardShaper::new(0.5, 0.0, 0.0);

        // Flat with signal => penalty applied
        let shaped = shaper.shape_reward(1.0_f32, true, true, 0);
        assert!((shaped - 0.5_f32).abs() < 1e-6, "Expected 0.5, got {}", shaped);

        // Flat without signal => no penalty
        shaper.reset();
        let shaped_no_signal = shaper.shape_reward(1.0_f32, true, false, 0);
        assert!(
            (shaped_no_signal - 1.0_f32).abs() < 1e-6,
            "Expected 1.0, got {}",
            shaped_no_signal
        );
    }

    #[test]
    fn test_reward_shaping_rolling_sharpe() {
        let mut shaper = PPORewardShaper::new(0.0, 1.0, 0.0);

        // Feed consistent positive returns to build up Sharpe
        for _ in 0..10 {
            shaper.shape_reward(1.0_f32, false, false, 0);
        }
        // All returns identical => std ~ 0, Sharpe returns 0 (division guard)
        // Feed slightly varied positive returns instead
        shaper.reset();
        for i in 0..10 {
            let reward = 1.0_f32 + (i as f32) * 0.1;
            shaper.shape_reward(reward, false, false, 0);
        }

        // Now check that Sharpe is positive (mean > 0, std > 0)
        let sharpe = shaper.current_sharpe();
        assert!(sharpe > 0.0, "Expected positive Sharpe, got {}", sharpe);

        // Shaped reward should be above raw because sharpe_weight * sharpe > 0
        let raw = 1.5_f32;
        let shaped = shaper.shape_reward(raw, false, false, 0);
        assert!(shaped > raw, "Expected shaped ({}) > raw ({})", shaped, raw);
    }

    #[test]
    fn test_reward_shaping_diversity_bonus() {
        // Diverse actions
        let mut diverse_shaper = PPORewardShaper::new(0.0, 0.0, 1.0);
        for i in 0..20 {
            diverse_shaper.shape_reward(0.0_f32, false, false, i % NUM_FACTORED_ACTIONS);
        }

        // Repeated single action
        let mut mono_shaper = PPORewardShaper::new(0.0, 0.0, 1.0);
        for _ in 0..20 {
            mono_shaper.shape_reward(0.0_f32, false, false, 0);
        }

        // The diverse shaper should yield higher cumulative diversity bonus.
        // We check the next call: diverse should give more bonus.
        let diverse_reward = diverse_shaper.shape_reward(0.0_f32, false, false, 5);
        let mono_reward = mono_shaper.shape_reward(0.0_f32, false, false, 0);
        assert!(
            diverse_reward > mono_reward,
            "Diverse ({}) should exceed monotone ({})",
            diverse_reward,
            mono_reward
        );
    }

    #[test]
    fn test_reward_shaping_reset() {
        let mut shaper = PPORewardShaper::new(0.0, 1.0, 1.0);

        // Accumulate some state
        for i in 0..15 {
            shaper.shape_reward(1.0_f32 + i as f32 * 0.1, false, false, i % NUM_FACTORED_ACTIONS);
        }
        assert!(!shaper.return_window.is_empty());
        assert!(shaper.total_actions > 0);

        // Reset clears everything
        shaper.reset();
        assert!(shaper.return_window.is_empty());
        assert_eq!(shaper.total_actions, 0);
        assert!(shaper.action_counts.iter().all(|&c| c == 0));
    }

    #[test]
    fn test_reward_shaping_default() {
        let shaper = PPORewardShaper::default();
        assert!((shaper.hold_penalty_weight() - 0.01).abs() < 1e-10);
        assert!((shaper.sharpe_weight() - 0.1).abs() < 1e-10);
        assert!((shaper.diversity_weight() - 0.05).abs() < 1e-10);
    }

    #[test]
    fn test_reward_shaping_all_components() {
        let mut shaper = PPORewardShaper::new(0.1, 0.5, 0.3);

        // Build up enough state for all components to fire
        for i in 0..15 {
            let r = 0.5_f32 + (i as f32) * 0.05;
            shaper.shape_reward(r, false, false, i % NUM_FACTORED_ACTIONS);
        }

        // Now shape a reward where all three components activate
        let raw = 1.0_f32;
        let shaped = shaper.shape_reward(raw, true, true, 3);

        // Hold penalty subtracts, Sharpe adds (positive returns), diversity adds
        // The result should differ from raw
        assert!(
            (shaped - raw).abs() > 1e-6,
            "All components should modify the reward, got {}",
            shaped
        );

        // Specifically: shaped < raw is possible because hold penalty (0.1) may dominate
        // but Sharpe + diversity add. Let's just check it's finite and reasonable.
        assert!(shaped.is_finite(), "Shaped reward should be finite");
    }

    #[test]
    fn test_reward_shaping_empty_window() {
        let mut shaper = PPORewardShaper::new(0.0, 1.0, 0.0);

        // With fewer than MIN_SHARPE_WINDOW returns, no Sharpe component
        for _ in 0..4 {
            let shaped = shaper.shape_reward(1.0_f32, false, false, 0);
            assert!(
                (shaped - 1.0_f32).abs() < 1e-6,
                "Expected no Sharpe with <5 returns, got {}",
                shaped
            );
        }

        // 5th return triggers Sharpe computation
        // But all returns are identical (1.0) => std ~ 0 => Sharpe returns 0
        let shaped_5th = shaper.shape_reward(2.0_f32, false, false, 0);
        // Returns are [1.0, 1.0, 1.0, 1.0, 2.0] — mean=1.2, std>0, Sharpe>0
        assert!(
            shaped_5th > 2.0,
            "Expected Sharpe bonus with 5+ returns, got {}",
            shaped_5th
        );
    }
}
