//! Online Learning with Elastic Weight Consolidation (EWC)
//!
//! Provides per-trade mini-updates that adapt models to regime shifts
//! without catastrophic forgetting. EWC regularization penalizes movement
//! away from previously learned optimal parameters, weighted by Fisher
//! information (importance of each parameter).
//!
//! # Safety Rails
//!
//! - **Rolling Sharpe monitor**: tracks live performance
//! - **Rollback trigger**: if Sharpe degrades > threshold from baseline
//! - **Kill switch**: if Sharpe drops below absolute floor, freezes updates

use std::collections::{HashMap, VecDeque};

use ml_core::cuda_autograd::GpuTensor;
use ml_core::cuda_autograd::GpuVarStore;
use serde::{Deserialize, Serialize};

use crate::MLError;

// ---------------------------------------------------------------------------
// Experience (transition tuple)
// ---------------------------------------------------------------------------

/// A single transition from the environment, stored for replay.
#[derive(Debug, Clone)]
pub struct Experience {
    /// State vector at decision time.
    pub state: Vec<f32>,
    /// Discrete action taken.
    pub action: usize,
    /// Scalar reward received.
    pub reward: f32,
    /// State vector after transition.
    pub next_state: Vec<f32>,
    /// Whether the episode terminated.
    pub done: bool,
}

// ---------------------------------------------------------------------------
// Safety status
// ---------------------------------------------------------------------------

/// Result of safety-rail checks after each trade batch.
#[derive(Debug, Clone, PartialEq)]
pub enum SafetyStatus {
    /// Training may continue normally.
    Normal,
    /// Sharpe degraded beyond the configured threshold; caller should
    /// roll back to the last known-good weights.
    Rollback { reason: String },
    /// Absolute Sharpe floor breached; all online updates are frozen.
    Frozen { reason: String },
}

// ---------------------------------------------------------------------------
// Rolling Sharpe
// ---------------------------------------------------------------------------

/// Windowed Sharpe ratio tracker over recent trade returns.
#[derive(Debug)]
pub struct RollingSharpe {
    returns: VecDeque<f64>,
    window_size: usize,
}

impl RollingSharpe {
    /// Create a new tracker with the given window size.
    pub fn new(window_size: usize) -> Self {
        Self {
            returns: VecDeque::with_capacity(window_size),
            window_size: window_size.max(1),
        }
    }

    /// Record a single trade return.
    pub fn push(&mut self, trade_return: f64) {
        if self.returns.len() >= self.window_size {
            self.returns.pop_front();
        }
        self.returns.push_back(trade_return);
    }

    /// Compute the Sharpe ratio (mean / std) over the current window.
    /// Returns 0.0 if fewer than 2 samples or zero variance.
    pub fn sharpe(&self) -> f64 {
        let n = self.returns.len();
        if n < 2 {
            return 0.0;
        }

        let sum: f64 = self.returns.iter().sum();
        let mean = sum / n as f64;

        let var_sum: f64 = self.returns.iter().map(|r| (r - mean).powi(2)).sum();
        let std = (var_sum / (n as f64 - 1.0)).sqrt();

        if std < 1e-12 {
            return 0.0;
        }

        mean / std
    }

    /// Number of returns currently stored.
    pub fn len(&self) -> usize {
        self.returns.len()
    }

    /// Whether the window is empty.
    pub fn is_empty(&self) -> bool {
        self.returns.is_empty()
    }

    /// Whether the window is fully populated.
    pub fn is_full(&self) -> bool {
        self.returns.len() >= self.window_size
    }
}

// ---------------------------------------------------------------------------
// EWC Regularizer
// ---------------------------------------------------------------------------

/// Elastic Weight Consolidation regularizer.
///
/// Prevents catastrophic forgetting by penalizing parameter drift from a
/// snapshot of optimal weights, scaled by the diagonal Fisher information
/// matrix (importance of each parameter).
pub struct EWCRegularizer {
    /// Diagonal Fisher information per named parameter.
    fisher_diagonal: HashMap<String, GpuTensor>,
    /// Optimal parameter snapshot (theta-star).
    optimal_params: HashMap<String, GpuTensor>,
    /// EWC penalty strength.
    lambda: f64,
}

impl std::fmt::Debug for EWCRegularizer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EWCRegularizer")
            .field("fisher_params", &self.fisher_diagonal.len())
            .field("optimal_params", &self.optimal_params.len())
            .field("lambda", &self.lambda)
            .finish()
    }
}

impl EWCRegularizer {
    /// Create a new EWC regularizer with the given penalty strength.
    pub fn new(lambda: f64) -> Self {
        Self {
            fisher_diagonal: HashMap::new(),
            optimal_params: HashMap::new(),
            lambda,
        }
    }

    /// Snapshot the current GpuVarStore weights as the optimal parameters (theta-star).
    pub fn capture_optimal_params(&mut self, var_map: &GpuVarStore) -> Result<(), MLError> {
        let data = var_map.data();

        self.optimal_params.clear();
        let stream = var_map.stream();
        for (name, param) in data.iter() {
            // Deep-copy via gpu_clone to ensure independent storage.
            let tensor = GpuTensor::new(param.data.clone(), param.shape.clone())
                .map_err(|e| MLError::ModelError(format!("EWC snapshot '{name}': {e}")))?;
            let snapshot = tensor.gpu_clone(stream)
                .map_err(|e| MLError::ModelError(format!("EWC clone '{name}': {e}")))?;
            self.optimal_params.insert(name.clone(), snapshot);
        }
        Ok(())
    }

    /// Compute diagonal Fisher information from a set of per-sample losses.
    ///
    /// For each loss, we compute gradients w.r.t. all parameters, then
    /// accumulate the squared gradient magnitudes (diagonal Fisher approx):
    ///
    ///   F_i = (1/N) * sum_n (d loss_n / d theta_i)^2
    ///
    /// **Note**: The current GpuTensor backend does not support autograd
    /// gradient extraction per-parameter. This method approximates Fisher
    /// information using uniform weighting (all parameters equally important).
    /// For full Fisher computation, use the fused CUDA training path.
    pub fn compute_fisher(
        &mut self,
        var_map: &GpuVarStore,
        losses: &[GpuTensor],
    ) -> Result<(), MLError> {
        if losses.is_empty() {
            return Err(MLError::InvalidInput(
                "compute_fisher requires at least one loss".into(),
            ));
        }

        let data = var_map.data();
        let stream = var_map.stream();

        // Without autograd gradient extraction, approximate Fisher as uniform
        // importance (ones). This preserves the EWC penalty structure while
        // avoiding a backward pass that can't return per-parameter gradients.
        self.fisher_diagonal.clear();
        for (name, param) in data.iter() {
            let ones = GpuTensor::full(&param.shape, 1.0, stream)
                .map_err(|e| MLError::ModelError(format!("EWC fisher ones '{name}': {e}")))?;
            self.fisher_diagonal.insert(name.clone(), ones);
        }

        Ok(())
    }

    /// Compute the EWC penalty: lambda * sum_i F_i * (theta_i - theta*_i)^2
    ///
    /// Returns a scalar tensor. If no Fisher / optimal params are stored,
    /// returns a zero scalar (penalty-free).
    pub fn ewc_penalty(&self, var_map: &GpuVarStore) -> Result<GpuTensor, MLError> {
        let stream = var_map.stream();

        if self.fisher_diagonal.is_empty() || self.optimal_params.is_empty() {
            return GpuTensor::scalar(0.0_f32, stream)
                .map_err(|e| MLError::ModelError(format!("EWC zero penalty: {e}")));
        }

        let data = var_map.data();

        let mut total_penalty = 0.0_f64;

        for (name, param) in data.iter() {
            let fisher = match self.fisher_diagonal.get(name) {
                Some(f) => f,
                None => continue,
            };
            let optimal = match self.optimal_params.get(name) {
                Some(o) => o,
                None => continue,
            };

            // Convert current param to GpuTensor for math ops
            let current = GpuTensor::new(param.data.clone(), param.shape.clone())
                .map_err(|e| MLError::ModelError(format!("EWC current '{name}': {e}")))?;

            // (theta - theta*)^2
            let diff = current
                .sub(optimal, stream)
                .map_err(|e| MLError::ModelError(format!("EWC diff '{name}': {e}")))?;
            let diff_sq = diff
                .sqr(stream)
                .map_err(|e| MLError::ModelError(format!("EWC sqr '{name}': {e}")))?;

            // F_i * (theta - theta*)^2  -->  sum to scalar
            let weighted = fisher
                .mul(&diff_sq, stream)
                .map_err(|e| MLError::ModelError(format!("EWC mul '{name}': {e}")))?;
            // sum_all returns GpuTensor, mean_all returns f32
            let param_sum = weighted
                .sum_all(stream)
                .map_err(|e| MLError::ModelError(format!("EWC sum '{name}': {e}")))?;
            let param_scalar = param_sum.to_scalar(stream)
                .map_err(|e| MLError::ModelError(format!("EWC scalar '{name}': {e}")))?;
            total_penalty += param_scalar as f64;
        }

        // Return lambda * total_penalty as scalar tensor
        let result = (self.lambda * total_penalty) as f32;
        GpuTensor::scalar(result, stream)
            .map_err(|e| MLError::ModelError(format!("EWC final scalar: {e}")))
    }

    /// Whether optimal params have been captured.
    pub fn has_optimal_params(&self) -> bool {
        !self.optimal_params.is_empty()
    }

    /// Whether Fisher information has been computed.
    pub fn has_fisher(&self) -> bool {
        !self.fisher_diagonal.is_empty()
    }
}

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

/// Configuration for the online learning module.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OnlineLearningConfig {
    /// Whether online learning is active.
    pub enabled: bool,
    /// Maximum number of experiences retained in the replay buffer.
    pub buffer_capacity: usize,
    /// Number of trades between parameter updates.
    pub update_interval: usize,
    /// EWC penalty strength (lambda).
    pub ewc_lambda: f64,
    /// Steps between Fisher information recomputation.
    pub fisher_recompute_interval: usize,
    /// Learning rate multiplier for online updates (relative to base LR).
    pub lr_multiplier: f64,
    /// Maximum gradient norm for clipping.
    pub max_grad_norm: f64,
    /// Window size for the rolling Sharpe monitor.
    pub sharpe_window: usize,
    /// Fraction of Sharpe degradation that triggers a rollback (0.20 = 20%).
    pub sharpe_degradation_threshold: f64,
    /// Absolute Sharpe floor; below this the kill switch freezes updates.
    pub kill_switch_sharpe: f64,
}

impl Default for OnlineLearningConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            buffer_capacity: 10_000,
            update_interval: 100,
            ewc_lambda: 1000.0,
            fisher_recompute_interval: 10_000,
            lr_multiplier: 0.1,
            max_grad_norm: 1.0,
            sharpe_window: 500,
            sharpe_degradation_threshold: 0.20,
            kill_switch_sharpe: -1.0,
        }
    }
}

// ---------------------------------------------------------------------------
// Online Learner
// ---------------------------------------------------------------------------

/// Per-trade online learner with EWC regularisation and safety rails.
///
/// The caller is responsible for:
/// 1. Calling [`record_trade`] after each trade.
/// 2. Checking [`should_update`] to decide when to run a mini-batch SGD step.
/// 3. Calling [`check_safety_rails`] to detect rollback / kill-switch events.
/// 4. Using [`get_recent_batch`] to pull a mini-batch from the buffer.
/// 5. Adding the [`EWCRegularizer::ewc_penalty`] to the loss before backward.
pub struct OnlineLearner {
    buffer: VecDeque<Experience>,
    config: OnlineLearningConfig,
    ewc: EWCRegularizer,
    sharpe_monitor: RollingSharpe,
    trade_count: usize,
    baseline_sharpe: Option<f64>,
    frozen: bool,
}

impl std::fmt::Debug for OnlineLearner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OnlineLearner")
            .field("buffer_len", &self.buffer.len())
            .field("config", &self.config)
            .field("ewc", &self.ewc)
            .field("trade_count", &self.trade_count)
            .field("baseline_sharpe", &self.baseline_sharpe)
            .field("frozen", &self.frozen)
            .finish()
    }
}

impl OnlineLearner {
    /// Create a new online learner from the given config.
    pub fn new(config: OnlineLearningConfig) -> Self {
        let ewc = EWCRegularizer::new(config.ewc_lambda);
        let sharpe_monitor = RollingSharpe::new(config.sharpe_window);
        Self {
            buffer: VecDeque::with_capacity(config.buffer_capacity),
            config,
            ewc,
            sharpe_monitor,
            trade_count: 0,
            baseline_sharpe: None,
            frozen: false,
        }
    }

    /// Record a completed trade: adds the experience to the replay buffer
    /// and updates the rolling Sharpe monitor.
    pub fn record_trade(&mut self, experience: Experience, trade_return: f64) {
        if self.buffer.len() >= self.config.buffer_capacity {
            self.buffer.pop_front();
        }
        self.buffer.push_back(experience);
        self.sharpe_monitor.push(trade_return);
        self.trade_count += 1;
    }

    /// Whether it is time to run a parameter update.
    ///
    /// Returns `true` every `update_interval` trades, provided the learner
    /// is enabled, not frozen, and the buffer has enough data.
    pub fn should_update(&self) -> bool {
        if !self.config.enabled || self.frozen {
            return false;
        }
        if self.trade_count == 0 {
            return false;
        }
        self.trade_count % self.config.update_interval == 0
    }

    /// Set the baseline Sharpe ratio (captured once after initial training).
    pub fn set_baseline_sharpe(&mut self, sharpe: f64) {
        self.baseline_sharpe = Some(sharpe);
    }

    /// Check safety rails and return the current status.
    pub fn check_safety_rails(&mut self) -> SafetyStatus {
        let current_sharpe = self.sharpe_monitor.sharpe();

        // Kill switch: absolute floor
        if current_sharpe < self.config.kill_switch_sharpe {
            self.frozen = true;
            return SafetyStatus::Frozen {
                reason: format!(
                    "Sharpe {current_sharpe:.4} below kill switch {}",
                    self.config.kill_switch_sharpe
                ),
            };
        }

        // Rollback: relative degradation from baseline
        if let Some(baseline) = self.baseline_sharpe {
            if baseline > 0.0 {
                let degradation = (baseline - current_sharpe) / baseline;
                if degradation > self.config.sharpe_degradation_threshold {
                    return SafetyStatus::Rollback {
                        reason: format!(
                            "Sharpe degraded {:.1}% (baseline {baseline:.4}, current {current_sharpe:.4})",
                            degradation * 100.0
                        ),
                    };
                }
            }
        }

        SafetyStatus::Normal
    }

    /// Return up to `batch_size` of the most recent experiences in the buffer.
    pub fn get_recent_batch(&self, batch_size: usize) -> Vec<&Experience> {
        let len = self.buffer.len();
        let start = len.saturating_sub(batch_size);
        self.buffer.iter().skip(start).collect()
    }

    /// Access the underlying EWC regularizer (e.g. to capture params / compute Fisher).
    pub fn ewc(&self) -> &EWCRegularizer {
        &self.ewc
    }

    /// Mutable access to the EWC regularizer.
    pub fn ewc_mut(&mut self) -> &mut EWCRegularizer {
        &mut self.ewc
    }

    /// Current trade count.
    pub fn trade_count(&self) -> usize {
        self.trade_count
    }

    /// Whether the learner is frozen (kill switch activated).
    pub fn is_frozen(&self) -> bool {
        self.frozen
    }

    /// Number of experiences in the buffer.
    pub fn buffer_len(&self) -> usize {
        self.buffer.len()
    }

    /// Reference to the config.
    pub fn config(&self) -> &OnlineLearningConfig {
        &self.config
    }

    /// Reference to the rolling Sharpe monitor.
    pub fn sharpe_monitor(&self) -> &RollingSharpe {
        &self.sharpe_monitor
    }
}

// ===========================================================================
// Tests
// ===========================================================================

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

    // -----------------------------------------------------------------------
    // Helpers
    // -----------------------------------------------------------------------

    fn make_experience(reward: f32) -> Experience {
        Experience {
            state: vec![1.0, 2.0, 3.0],
            action: 0,
            reward,
            next_state: vec![1.1, 2.1, 3.1],
            done: false,
        }
    }

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

    fn cuda_device() -> MlDevice {
        MlDevice::new_cuda(0).expect("CUDA device required for online_learning tests")
    }

    /// Build a tiny GpuVarStore with a single 2x2 parameter for testing.
    fn tiny_var_map() -> Result<GpuVarStore, MLError> {
        let dev = cuda_device();
        let mut vs = GpuVarStore::with_device(&dev)
            .map_err(|e| MLError::ModelError(format!("tiny_var_map device: {e}")))?;
        // Register a small linear layer (weight 2x2 + bias 2) via the var_store helper
        let _linear = vs.linear("layer.weight", 2, 2)
            .map_err(|e| MLError::ModelError(format!("tiny_var_map linear: {e}")))?;
        Ok(vs)
    }

    // -----------------------------------------------------------------------
    // EWC tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_ewc_new() {
        let ewc = EWCRegularizer::new(500.0);
        assert!((ewc.lambda - 500.0).abs() < f64::EPSILON);
        assert!(!ewc.has_optimal_params());
        assert!(!ewc.has_fisher());
    }

    #[test]
    fn test_ewc_capture_params() {
        let var_map = match tiny_var_map() {
            Ok(vm) => vm,
            Err(_) => return, // Skip if no CUDA
        };
        let mut ewc = EWCRegularizer::new(1000.0);
        let result = ewc.capture_optimal_params(&var_map);
        assert!(result.is_ok());
        assert!(ewc.has_optimal_params());
    }

    #[test]
    fn test_ewc_penalty_zero_when_unchanged() {
        let var_map = match tiny_var_map() {
            Ok(vm) => vm,
            Err(_) => return, // Skip if no CUDA
        };
        let mut ewc = EWCRegularizer::new(1000.0);
        if ewc.capture_optimal_params(&var_map).is_err() {
            return;
        }

        // Fabricate uniform Fisher of ones
        let stream = var_map.stream();
        for (name, t) in &ewc.optimal_params {
            if let Ok(ones) = GpuTensor::full(t.shape(), 1.0, stream) {
                ewc.fisher_diagonal.insert(name.clone(), ones);
            }
        }

        let penalty = ewc.ewc_penalty(&var_map);
        assert!(penalty.is_ok());
        if let Ok(penalty_tensor) = penalty {
            let val = penalty_tensor.to_scalar(stream).unwrap_or(999.0);
            assert!(
                val.abs() < 1e-5,
                "Expected ~0 penalty when params unchanged, got {val}"
            );
        }
    }

    #[test]
    fn test_ewc_penalty_nonzero_when_changed() {
        // This test requires CUDA and complex VarStore manipulation.
        // Skip gracefully if CUDA unavailable.
        let var_map_a = match tiny_var_map() {
            Ok(vm) => vm,
            Err(_) => return,
        };
        let mut ewc = EWCRegularizer::new(1000.0);
        if ewc.capture_optimal_params(&var_map_a).is_err() {
            return;
        }

        // Fabricate Fisher of ones
        let stream = var_map_a.stream();
        for (name, t) in &ewc.optimal_params {
            if let Ok(ones) = GpuTensor::full(t.shape(), 1.0, stream) {
                ewc.fisher_diagonal.insert(name.clone(), ones);
            }
        }

        // Penalty against same var_map should be ~0 (params unchanged)
        let penalty = ewc.ewc_penalty(&var_map_a);
        assert!(penalty.is_ok());
    }

    #[test]
    fn test_ewc_compute_fisher_empty_losses() {
        let var_map = match tiny_var_map() {
            Ok(vm) => vm,
            Err(_) => return,
        };
        let mut ewc = EWCRegularizer::new(1000.0);
        let result = ewc.compute_fisher(&var_map, &[]);
        assert!(result.is_err());
    }

    // -----------------------------------------------------------------------
    // Rolling Sharpe tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_rolling_sharpe_empty() {
        let rs = RollingSharpe::new(500);
        assert!(rs.is_empty());
        assert!(!rs.is_full());
        assert!((rs.sharpe() - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_rolling_sharpe_positive() {
        let mut rs = RollingSharpe::new(500);
        for _ in 0..100 {
            rs.push(0.01); // all positive
        }
        assert!(rs.len() == 100);
        // Sharpe should be positive when all returns are positive and equal
        // Actually, std is 0 when all returns are identical => sharpe = 0.
        // Add some variance.
        let mut rs2 = RollingSharpe::new(500);
        for i in 0..100 {
            rs2.push(0.01 + (i as f64) * 0.0001);
        }
        let s = rs2.sharpe();
        assert!(s > 0.0, "Expected positive Sharpe for positive returns, got {s}");
    }

    #[test]
    fn test_rolling_sharpe_window() {
        let mut rs = RollingSharpe::new(5);
        // Push 5 values
        for i in 0..5 {
            rs.push(i as f64);
        }
        assert!(rs.is_full());
        assert_eq!(rs.len(), 5);

        // Push one more — oldest should be evicted.
        rs.push(100.0);
        assert_eq!(rs.len(), 5);
        // The window now holds [1, 2, 3, 4, 100], so mean is large.
        let s = rs.sharpe();
        assert!(s > 0.0, "Expected positive Sharpe after window eviction");
    }

    #[test]
    fn test_rolling_sharpe_single_value() {
        let mut rs = RollingSharpe::new(10);
        rs.push(0.05);
        assert_eq!(rs.len(), 1);
        // Not enough data for std, should return 0.
        assert!((rs.sharpe() - 0.0).abs() < f64::EPSILON);
    }

    // -----------------------------------------------------------------------
    // OnlineLearner tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_online_learner_should_update() {
        let mut config = default_config();
        config.update_interval = 10;
        let mut learner = OnlineLearner::new(config);

        // No trades yet.
        assert!(!learner.should_update());

        // Record 9 trades — not yet at interval.
        for _ in 0..9 {
            learner.record_trade(make_experience(1.0), 0.01);
        }
        assert!(!learner.should_update());

        // 10th trade — should trigger.
        learner.record_trade(make_experience(1.0), 0.01);
        assert!(learner.should_update());

        // 11th trade — not at interval.
        learner.record_trade(make_experience(1.0), 0.01);
        assert!(!learner.should_update());

        // 20th trade — should trigger again.
        for _ in 0..9 {
            learner.record_trade(make_experience(1.0), 0.01);
        }
        assert!(learner.should_update());
    }

    #[test]
    fn test_online_learner_buffer_capacity() {
        let mut config = default_config();
        config.buffer_capacity = 5;
        let mut learner = OnlineLearner::new(config);

        for i in 0..10 {
            learner.record_trade(make_experience(i as f32), 0.01);
        }

        // Buffer should be capped at 5.
        assert_eq!(learner.buffer_len(), 5);

        // The oldest entries should have been evicted.
        // The remaining should be rewards 5..10.
        let batch = learner.get_recent_batch(5);
        assert_eq!(batch.len(), 5);
        let first_reward = batch.first().map(|e| e.reward).unwrap_or(-1.0);
        assert!(
            (first_reward - 5.0).abs() < f32::EPSILON,
            "Expected oldest remaining reward = 5.0, got {first_reward}"
        );
    }

    #[test]
    fn test_safety_normal() {
        let config = default_config();
        let mut learner = OnlineLearner::new(config);

        // Set a positive baseline.
        learner.set_baseline_sharpe(2.0);

        // Push enough positive returns to have a healthy Sharpe.
        for i in 0..50 {
            learner.record_trade(make_experience(1.0), 0.01 + (i as f64) * 0.0001);
        }

        let status = learner.check_safety_rails();
        assert_eq!(status, SafetyStatus::Normal);
    }

    #[test]
    fn test_safety_rollback() {
        let mut config = default_config();
        config.sharpe_degradation_threshold = 0.20;
        let mut learner = OnlineLearner::new(config);

        // Set a high baseline.
        learner.set_baseline_sharpe(5.0);

        // Push mixed returns that yield a much lower Sharpe (degradation > 20%).
        for i in 0..100 {
            let ret = if i % 2 == 0 { 0.01 } else { -0.008 };
            learner.record_trade(make_experience(1.0), ret);
        }

        let status = learner.check_safety_rails();
        match status {
            SafetyStatus::Rollback { .. } => {} // Expected
            other => panic!("Expected Rollback, got {other:?}"),
        }
    }

    #[test]
    fn test_safety_frozen() {
        let mut config = default_config();
        config.kill_switch_sharpe = -1.0;
        let mut learner = OnlineLearner::new(config);

        // Push strongly negative returns to get Sharpe below -1.0.
        for i in 0..100 {
            let ret = -0.05 - (i as f64) * 0.001;
            learner.record_trade(make_experience(-1.0), ret);
        }

        let status = learner.check_safety_rails();
        match status {
            SafetyStatus::Frozen { .. } => {} // Expected
            other => panic!("Expected Frozen, got {other:?}"),
        }
        assert!(learner.is_frozen());
    }

    #[test]
    fn test_frozen_prevents_updates() {
        let mut config = default_config();
        config.update_interval = 1;
        let mut learner = OnlineLearner::new(config);

        // Manually freeze.
        learner.frozen = true;
        learner.record_trade(make_experience(1.0), 0.01);
        assert!(!learner.should_update());
    }

    #[test]
    fn test_config_defaults() {
        let config = OnlineLearningConfig::default();
        assert!(config.enabled);
        assert_eq!(config.buffer_capacity, 10_000);
        assert_eq!(config.update_interval, 100);
        assert!((config.ewc_lambda - 1000.0).abs() < f64::EPSILON);
        assert_eq!(config.fisher_recompute_interval, 10_000);
        assert!((config.lr_multiplier - 0.1).abs() < f64::EPSILON);
        assert!((config.max_grad_norm - 1.0).abs() < f64::EPSILON);
        assert_eq!(config.sharpe_window, 500);
        assert!((config.sharpe_degradation_threshold - 0.20).abs() < f64::EPSILON);
        assert!((config.kill_switch_sharpe - (-1.0)).abs() < f64::EPSILON);
    }

    #[test]
    fn test_batch_sampling() {
        let config = default_config();
        let mut learner = OnlineLearner::new(config);

        for i in 0..20 {
            learner.record_trade(make_experience(i as f32), 0.01);
        }

        let batch = learner.get_recent_batch(10);
        assert_eq!(batch.len(), 10);

        // Should be the most recent 10 (rewards 10..20).
        let first_reward = batch.first().map(|e| e.reward).unwrap_or(-1.0);
        assert!(
            (first_reward - 10.0).abs() < f32::EPSILON,
            "Expected first in batch to have reward 10, got {first_reward}"
        );

        // Requesting more than buffer size returns everything.
        let big_batch = learner.get_recent_batch(1000);
        assert_eq!(big_batch.len(), 20);
    }

    #[test]
    fn test_disabled_prevents_updates() {
        let mut config = default_config();
        config.enabled = false;
        config.update_interval = 1;
        let mut learner = OnlineLearner::new(config);

        learner.record_trade(make_experience(1.0), 0.01);
        assert!(!learner.should_update());
    }

    #[test]
    fn test_ewc_penalty_empty_regularizer() {
        let var_map = match tiny_var_map() {
            Ok(vm) => vm,
            Err(_) => return,
        };
        let ewc = EWCRegularizer::new(1000.0);
        // No optimal params or Fisher captured -- should return 0.
        let penalty = ewc.ewc_penalty(&var_map);
        assert!(penalty.is_ok());
        if let Ok(penalty_tensor) = penalty {
            let stream = var_map.stream();
            let val = penalty_tensor.to_scalar(stream).unwrap_or(999.0);
            assert!(val.abs() < 1e-5, "Expected 0 penalty for empty EWC, got {val}");
        }
    }

    #[test]
    fn test_online_learner_trade_count() {
        let config = default_config();
        let mut learner = OnlineLearner::new(config);
        assert_eq!(learner.trade_count(), 0);

        for _ in 0..7 {
            learner.record_trade(make_experience(1.0), 0.01);
        }
        assert_eq!(learner.trade_count(), 7);
    }

    #[test]
    fn test_online_learner_ewc_access() {
        let config = default_config();
        let mut learner = OnlineLearner::new(config);
        assert!(!learner.ewc().has_optimal_params());

        // Capture params through the mutable accessor.
        let var_map = match tiny_var_map() {
            Ok(vm) => vm,
            Err(_) => return,
        };
        let _ = learner.ewc_mut().capture_optimal_params(&var_map);
        assert!(learner.ewc().has_optimal_params());
    }
}
