//! Curriculum Learning for RL Agent Training
//!
//! Phases in action complexity gradually to accelerate convergence:
//!
//! | Phase | Exposure Actions | Count | Regimes | Gate |
//! |-------|-----------------|-------|---------|------|
//! | 1: Basic | Short50, Flat, Long50 | 3 | Ranging only | Sharpe > 0.5 for 3 folds |
//! | 2: FullPosition | All 5 exposures | 5 | Ranging + Trending | Sharpe > 0.3 for 2 folds |
//! | 3: AllRegimes | All 5 | 5 | All regimes | Terminal |
//!
//! The scheduler tracks consecutive qualifying folds and advances the phase
//! when the gate condition is met. Action masks combine curriculum restrictions
//! with position-limit masks from [`crate::dqn::action_space`].

use serde::{Deserialize, Serialize};

use crate::dqn::action_space::get_valid_action_mask;
use crate::dqn::regime_conditional::RegimeType;

/// Number of DQN exposure actions (Short100, Short50, Flat, Long50, Long100).
const NUM_ACTIONS: usize = 5;

/// Curriculum training phase.
///
/// Phases are ordered by complexity; the agent must clear a Sharpe gate
/// before advancing to the next phase.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CurriculumPhase {
    /// Phase 1: limited exposures (Short50, Flat, Long50), Ranging regime only.
    Basic,
    /// Phase 2: all 5 exposures, Ranging + Trending regimes.
    FullPosition,
    /// Phase 3: all actions, all regimes (terminal phase).
    AllRegimes,
}

/// Configuration for the curriculum scheduler.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CurriculumConfig {
    /// Whether curriculum learning is enabled.
    pub enabled: bool,
    /// Sharpe threshold to advance from Phase 1 to Phase 2.
    pub phase1_sharpe_threshold: f64,
    /// Number of consecutive folds meeting the Phase 1 threshold.
    pub phase1_consecutive_folds: usize,
    /// Sharpe threshold to advance from Phase 2 to Phase 3.
    pub phase2_sharpe_threshold: f64,
    /// Number of consecutive folds meeting the Phase 2 threshold.
    pub phase2_consecutive_folds: usize,
    /// Manual override: skip directly to this phase.
    pub override_phase: Option<CurriculumPhase>,
}

impl Default for CurriculumConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            phase1_sharpe_threshold: 0.5,
            phase1_consecutive_folds: 3,
            phase2_sharpe_threshold: 0.3,
            phase2_consecutive_folds: 2,
            override_phase: None,
        }
    }
}

/// Curriculum scheduler that manages phase progression and action masking.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CurriculumScheduler {
    current_phase: CurriculumPhase,
    config: CurriculumConfig,
    consecutive_qualifying_folds: usize,
}

impl CurriculumScheduler {
    /// Create a new scheduler.
    ///
    /// Precedence: `!enabled` (→ AllRegimes) > `override_phase` > default (Basic).
    /// When disabled, `override_phase` is ignored — all actions/regimes are allowed.
    pub fn new(config: CurriculumConfig) -> Self {
        let initial_phase = if !config.enabled {
            CurriculumPhase::AllRegimes
        } else if let Some(phase) = config.override_phase {
            phase
        } else {
            CurriculumPhase::Basic
        };

        Self {
            current_phase: initial_phase,
            config,
            consecutive_qualifying_folds: 0,
        }
    }

    /// Return the current curriculum phase.
    pub fn current_phase(&self) -> CurriculumPhase {
        self.current_phase
    }

    /// Record the Sharpe ratio from a completed fold.
    ///
    /// Returns `true` if the phase advanced as a result.
    pub fn record_fold_result(&mut self, sharpe: f64) -> bool {
        match self.current_phase {
            CurriculumPhase::Basic => {
                if sharpe >= self.config.phase1_sharpe_threshold {
                    self.consecutive_qualifying_folds += 1;
                    if self.consecutive_qualifying_folds >= self.config.phase1_consecutive_folds {
                        self.current_phase = CurriculumPhase::FullPosition;
                        self.consecutive_qualifying_folds = 0;
                        return true;
                    }
                } else {
                    self.consecutive_qualifying_folds = 0;
                }
                false
            }
            CurriculumPhase::FullPosition => {
                if sharpe >= self.config.phase2_sharpe_threshold {
                    self.consecutive_qualifying_folds += 1;
                    if self.consecutive_qualifying_folds >= self.config.phase2_consecutive_folds {
                        self.current_phase = CurriculumPhase::AllRegimes;
                        self.consecutive_qualifying_folds = 0;
                        return true;
                    }
                } else {
                    self.consecutive_qualifying_folds = 0;
                }
                false
            }
            CurriculumPhase::AllRegimes => {
                // Terminal phase; never advances.
                false
            }
        }
    }

    /// Generate a 5-element exposure action mask for the current phase.
    ///
    /// The mask combines:
    /// 1. Curriculum restriction (phase-dependent exposure filter)
    /// 2. Position-limit mask from `get_valid_action_mask`
    ///
    /// An exposure action is valid only if both masks permit it.
    pub fn get_action_mask(&self, max_position: f64) -> Vec<bool> {
        let position_mask = get_valid_action_mask(0.0, max_position);
        let curriculum_mask = self.curriculum_phase_mask();

        // Element-wise AND of both masks.
        position_mask
            .iter()
            .zip(curriculum_mask.iter())
            .map(|(&pos, &cur)| pos && cur)
            .collect()
    }

    /// Return the regimes allowed for training data in the current phase.
    pub fn allowed_regimes(&self) -> Vec<RegimeType> {
        match self.current_phase {
            CurriculumPhase::Basic => {
                // Ranging is the "normal" / default regime.
                vec![RegimeType::Ranging]
            }
            CurriculumPhase::FullPosition => {
                vec![RegimeType::Ranging, RegimeType::Trending]
            }
            CurriculumPhase::AllRegimes => {
                vec![RegimeType::Ranging, RegimeType::Trending, RegimeType::Volatile]
            }
        }
    }

    /// Build a raw 5-element mask based purely on the curriculum phase
    /// (no position-limit filtering).
    fn curriculum_phase_mask(&self) -> Vec<bool> {
        match self.current_phase {
            CurriculumPhase::Basic => self.phase1_mask(),
            CurriculumPhase::FullPosition | CurriculumPhase::AllRegimes => {
                // Phase 2 and 3 allow all 5 exposure actions.
                vec![true; NUM_ACTIONS]
            }
        }
    }

    /// Phase 1 mask: only Short50 (idx 1), Flat (idx 2), Long50 (idx 3) exposures.
    /// Masks out Short100 (idx 0) and Long100 (idx 4).
    fn phase1_mask(&self) -> Vec<bool> {
        // [Short100, Short50, Flat, Long50, Long100]
        vec![false, true, true, true, false]
    }
}

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

    fn default_config() -> CurriculumConfig {
        CurriculumConfig::default()
    }

    // ---------------------------------------------------------------
    // Test 1: Phase 1 masks extreme exposures (Short100, Long100)
    // ---------------------------------------------------------------
    #[test]
    fn test_phase1_masks_extreme_exposure() {
        let scheduler = CurriculumScheduler::new(default_config());
        assert_eq!(scheduler.current_phase(), CurriculumPhase::Basic);

        let mask = scheduler.get_action_mask(2.0);
        assert_eq!(mask.len(), NUM_ACTIONS);

        // Short100 (idx 0) masked
        assert!(!mask[0], "Short100 should be masked in Phase 1");
        // Long100 (idx 4) masked
        assert!(!mask[4], "Long100 should be masked in Phase 1");
    }

    // ---------------------------------------------------------------
    // Test 2: Phase 1 allows Short50, Flat, Long50 (3 actions)
    // ---------------------------------------------------------------
    #[test]
    fn test_phase1_allows_core_actions() {
        let scheduler = CurriculumScheduler::new(default_config());
        let mask = scheduler.get_action_mask(2.0);

        let allowed_count = mask.iter().filter(|&&v| v).count();
        assert_eq!(
            allowed_count, 3,
            "Phase 1 should allow exactly 3 exposure actions (Short50, Flat, Long50)"
        );

        assert!(mask[1], "Short50 should be allowed in Phase 1");
        assert!(mask[2], "Flat should be allowed in Phase 1");
        assert!(mask[3], "Long50 should be allowed in Phase 1");
    }

    // ---------------------------------------------------------------
    // Test 3: Phase 2 allows all 5 exposure actions
    // ---------------------------------------------------------------
    #[test]
    fn test_phase2_allows_all_exposures() {
        let mut scheduler = CurriculumScheduler::new(default_config());
        // Advance to Phase 2
        scheduler.record_fold_result(0.6);
        scheduler.record_fold_result(0.6);
        assert!(scheduler.record_fold_result(0.6));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);

        let mask = scheduler.get_action_mask(2.0);
        let allowed = mask.iter().filter(|&&v| v).count();
        assert_eq!(allowed, 5, "Phase 2 should allow all 5 exposure actions");
    }

    // ---------------------------------------------------------------
    // Test 4: Phase 3 has same mask as Phase 2
    // ---------------------------------------------------------------
    #[test]
    fn test_phase3_same_as_phase2() {
        let mut scheduler = CurriculumScheduler::new(default_config());
        // Advance to Phase 2
        for _ in 0..3 {
            scheduler.record_fold_result(0.6);
        }
        let phase2_mask = scheduler.get_action_mask(2.0);

        // Advance to Phase 3
        for _ in 0..2 {
            scheduler.record_fold_result(0.4);
        }
        assert_eq!(scheduler.current_phase(), CurriculumPhase::AllRegimes);

        let phase3_mask = scheduler.get_action_mask(2.0);
        assert_eq!(phase2_mask, phase3_mask, "Phase 3 mask should equal Phase 2 mask");
    }

    // ---------------------------------------------------------------
    // Test 5: Sharpe advancement Phase 1 -> Phase 2
    // ---------------------------------------------------------------
    #[test]
    fn test_sharpe_advancement_phase1_to_2() {
        let mut scheduler = CurriculumScheduler::new(default_config());
        assert_eq!(scheduler.current_phase(), CurriculumPhase::Basic);

        // Two qualifying folds -- not enough
        assert!(!scheduler.record_fold_result(0.5));
        assert!(!scheduler.record_fold_result(0.7));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::Basic);

        // Third consecutive qualifying fold -> advance
        assert!(scheduler.record_fold_result(0.5));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);
    }

    // ---------------------------------------------------------------
    // Test 6: Sharpe advancement Phase 2 -> Phase 3
    // ---------------------------------------------------------------
    #[test]
    fn test_sharpe_advancement_phase2_to_3() {
        let mut scheduler = CurriculumScheduler::new(default_config());
        // Advance to Phase 2 first
        for _ in 0..3 {
            scheduler.record_fold_result(0.6);
        }
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);

        // One qualifying fold -- not enough
        assert!(!scheduler.record_fold_result(0.3));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);

        // Second consecutive qualifying fold -> advance
        assert!(scheduler.record_fold_result(0.4));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::AllRegimes);
    }

    // ---------------------------------------------------------------
    // Test 7: Non-consecutive fold resets counter
    // ---------------------------------------------------------------
    #[test]
    fn test_non_consecutive_resets_counter() {
        let mut scheduler = CurriculumScheduler::new(default_config());

        // Two qualifying folds, then a failure
        scheduler.record_fold_result(0.6);
        scheduler.record_fold_result(0.6);
        scheduler.record_fold_result(0.1); // Reset!

        // Need 3 more consecutive qualifying folds now
        assert!(!scheduler.record_fold_result(0.5));
        assert!(!scheduler.record_fold_result(0.5));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::Basic);
        assert!(scheduler.record_fold_result(0.5)); // 3rd consecutive -> advance
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);
    }

    // ---------------------------------------------------------------
    // Test 8: Phase 3 never advances
    // ---------------------------------------------------------------
    #[test]
    fn test_phase3_never_advances() {
        let mut scheduler = CurriculumScheduler::new(default_config());
        // Advance to Phase 3
        for _ in 0..3 {
            scheduler.record_fold_result(0.6);
        }
        for _ in 0..2 {
            scheduler.record_fold_result(0.4);
        }
        assert_eq!(scheduler.current_phase(), CurriculumPhase::AllRegimes);

        // Many high-Sharpe folds should not change anything
        for _ in 0..20 {
            assert!(!scheduler.record_fold_result(10.0));
        }
        assert_eq!(scheduler.current_phase(), CurriculumPhase::AllRegimes);
    }

    // ---------------------------------------------------------------
    // Test 9: Override phase skips to specified phase
    // ---------------------------------------------------------------
    #[test]
    fn test_override_phase() {
        let config = CurriculumConfig {
            override_phase: Some(CurriculumPhase::FullPosition),
            ..Default::default()
        };
        let scheduler = CurriculumScheduler::new(config);
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);

        // AllRegimes override
        let config2 = CurriculumConfig {
            override_phase: Some(CurriculumPhase::AllRegimes),
            ..Default::default()
        };
        let scheduler2 = CurriculumScheduler::new(config2);
        assert_eq!(scheduler2.current_phase(), CurriculumPhase::AllRegimes);
    }

    // ---------------------------------------------------------------
    // Test 10: Correct regime filtering per phase
    // ---------------------------------------------------------------
    #[test]
    fn test_allowed_regimes_per_phase() {
        let mut scheduler = CurriculumScheduler::new(default_config());

        // Phase 1: Ranging only
        let regimes1 = scheduler.allowed_regimes();
        assert_eq!(regimes1, vec![RegimeType::Ranging]);

        // Phase 2: Ranging + Trending
        for _ in 0..3 {
            scheduler.record_fold_result(0.6);
        }
        let regimes2 = scheduler.allowed_regimes();
        assert_eq!(regimes2.len(), 2);
        assert!(regimes2.contains(&RegimeType::Ranging));
        assert!(regimes2.contains(&RegimeType::Trending));

        // Phase 3: All 3 regimes
        for _ in 0..2 {
            scheduler.record_fold_result(0.4);
        }
        let regimes3 = scheduler.allowed_regimes();
        assert_eq!(regimes3.len(), 3);
        assert!(regimes3.contains(&RegimeType::Ranging));
        assert!(regimes3.contains(&RegimeType::Trending));
        assert!(regimes3.contains(&RegimeType::Volatile));
    }

    // ---------------------------------------------------------------
    // Test 11: Default config is sensible
    // ---------------------------------------------------------------
    #[test]
    fn test_config_defaults() {
        let config = CurriculumConfig::default();
        assert!(config.enabled);
        assert!((config.phase1_sharpe_threshold - 0.5).abs() < f64::EPSILON);
        assert_eq!(config.phase1_consecutive_folds, 3);
        assert!((config.phase2_sharpe_threshold - 0.3).abs() < f64::EPSILON);
        assert_eq!(config.phase2_consecutive_folds, 2);
        assert!(config.override_phase.is_none());
    }

    // ---------------------------------------------------------------
    // Test 12: Curriculum disabled -> all actions from start
    // ---------------------------------------------------------------
    #[test]
    fn test_curriculum_disabled() {
        let config = CurriculumConfig {
            enabled: false,
            ..Default::default()
        };
        let scheduler = CurriculumScheduler::new(config);
        assert_eq!(scheduler.current_phase(), CurriculumPhase::AllRegimes);

        let mask = scheduler.get_action_mask(2.0);
        let allowed = mask.iter().filter(|&&v| v).count();
        assert_eq!(allowed, 5, "Disabled curriculum should allow all 5 exposure actions");

        let regimes = scheduler.allowed_regimes();
        assert_eq!(regimes.len(), 3, "Disabled curriculum should allow all regimes");
    }

    // ---------------------------------------------------------------
    // Extra: position-limit interaction
    // ---------------------------------------------------------------
    #[test]
    fn test_position_limit_interaction_with_curriculum() {
        // With max_position=0.6, Short100/Long100 masked by position limit,
        // AND Phase 1 also masks them. Combined: only Short50, Flat, Long50.
        let scheduler = CurriculumScheduler::new(default_config());
        let mask = scheduler.get_action_mask(0.6);

        assert!(!mask[0], "Short100 masked by both curriculum and position limit");
        assert!(mask[1], "Short50 allowed (Phase 1 + within 0.6 limit)");
        assert!(mask[2], "Flat allowed");
        assert!(mask[3], "Long50 allowed (Phase 1 + within 0.6 limit)");
        assert!(!mask[4], "Long100 masked by both curriculum and position limit");
    }

    // ---------------------------------------------------------------
    // Extra: Phase 2 with restrictive position limit
    // ---------------------------------------------------------------
    #[test]
    fn test_phase2_with_restrictive_position_limit() {
        let mut scheduler = CurriculumScheduler::new(default_config());
        // Advance to Phase 2
        for _ in 0..3 {
            scheduler.record_fold_result(0.6);
        }
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);

        // Position limit 0.6 masks Short100/Long100 even though curriculum allows them
        let mask = scheduler.get_action_mask(0.6);

        assert!(!mask[0], "Short100 masked by position limit");
        assert!(mask[1], "Short50 allowed");
        assert!(mask[2], "Flat allowed");
        assert!(mask[3], "Long50 allowed");
        assert!(!mask[4], "Long100 masked by position limit");
    }

    // ---------------------------------------------------------------
    // Extra: Boundary Sharpe exactly at threshold
    // ---------------------------------------------------------------
    #[test]
    fn test_sharpe_exactly_at_threshold() {
        let mut scheduler = CurriculumScheduler::new(default_config());

        // Sharpe exactly at 0.5 should qualify (>=)
        for _ in 0..3 {
            scheduler.record_fold_result(0.5);
        }
        assert_eq!(
            scheduler.current_phase(),
            CurriculumPhase::FullPosition,
            "Sharpe exactly at threshold should qualify"
        );
    }

    // ---------------------------------------------------------------
    // Extra: Sharpe just below threshold does not qualify
    // ---------------------------------------------------------------
    #[test]
    fn test_sharpe_below_threshold_no_advancement() {
        let mut scheduler = CurriculumScheduler::new(default_config());

        // Sharpe just below 0.5 should NOT qualify
        for _ in 0..100 {
            scheduler.record_fold_result(0.4999);
        }
        assert_eq!(
            scheduler.current_phase(),
            CurriculumPhase::Basic,
            "Sharpe below threshold should not advance"
        );
    }

    // ---------------------------------------------------------------
    // Extra: Override to Basic starts correctly
    // ---------------------------------------------------------------
    #[test]
    fn test_override_to_basic() {
        let config = CurriculumConfig {
            override_phase: Some(CurriculumPhase::Basic),
            ..Default::default()
        };
        let scheduler = CurriculumScheduler::new(config);
        assert_eq!(scheduler.current_phase(), CurriculumPhase::Basic);
    }

    // ---------------------------------------------------------------
    // Extra: Disabled config override_phase is ignored
    // ---------------------------------------------------------------
    #[test]
    fn test_disabled_ignores_override() {
        let config = CurriculumConfig {
            enabled: false,
            override_phase: Some(CurriculumPhase::Basic),
            ..Default::default()
        };
        let scheduler = CurriculumScheduler::new(config);
        // disabled takes precedence
        assert_eq!(scheduler.current_phase(), CurriculumPhase::AllRegimes);
    }

    // ---------------------------------------------------------------
    // Extra: Consecutive counter resets on phase advance
    // ---------------------------------------------------------------
    #[test]
    fn test_counter_resets_on_advance() {
        let mut scheduler = CurriculumScheduler::new(default_config());

        // Advance Phase 1 -> Phase 2 with 3 qualifying folds
        for _ in 0..3 {
            scheduler.record_fold_result(0.6);
        }
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);

        // Counter should have reset; a single qualifying fold at phase2 threshold
        // should NOT advance yet (need 2 consecutive)
        assert!(!scheduler.record_fold_result(0.4));
        assert_eq!(scheduler.current_phase(), CurriculumPhase::FullPosition);
    }
}
