// Bug #19 + Bug #20 Integration Test
// Tests that both fixes work together to prevent gradient collapse

use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use ml::dqn::portfolio_tracker::PortfolioTracker;
use ml::dqn::{WorkingDQN, WorkingDQNConfig};

#[test]
fn test_gradient_stability_with_normalized_portfolio() -> Result<()> {
    // Test gradient stability across multiple epochs with normalized portfolio features
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.num_actions = 45;
    config.hidden_dims = vec![512, 256];

    let dqn = WorkingDQN::new(config)?;
    let tracker = PortfolioTracker::new(100_000.0, 0.0001, 0.0);
    let device = dqn.device().clone();

    // Simulate 5 epochs of 100 steps each
    for epoch in 0..5 {
        let mut epoch_q_values = Vec::new();

        for step in 0..100 {
            // Create state with normalized portfolio
            let mut state_vec = vec![0.0f32; 128];
            let portfolio_features = tracker.get_portfolio_features(4000.0);
            state_vec[0..3].copy_from_slice(&portfolio_features);

            // Add market features
            for i in 3..128 {
                state_vec[i] = ((step + epoch * 100) as f32 * 0.01).sin() * 0.1;
            }

            let state = Tensor::from_vec(state_vec, (1, 128), &device)?.to_dtype(DType::F32)?;
            let q_values = dqn.forward(&state)?;
            let q_vec = q_values.flatten_all()?.to_vec1::<f32>()?;

            epoch_q_values.extend_from_slice(&q_vec);
        }

        // Check Q-values are finite
        let max_q = epoch_q_values
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);
        let min_q = epoch_q_values
            .iter()
            .copied()
            .fold(f32::INFINITY, f32::min);

        println!("Epoch {}: Q-range=[{:.2}, {:.2}]", epoch, min_q, max_q);

        assert!(
            max_q.is_finite() && min_q.is_finite(),
            "Q-values became NaN/Inf in epoch {}",
            epoch
        );
    }

    println!("✅ Gradient stability maintained across all epochs");

    Ok(())
}

#[test]
fn test_no_q_value_explosion_with_normalized_portfolio() -> Result<()> {
    // Test that normalized portfolio features prevent Q-value explosion
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.num_actions = 45;
    config.hidden_dims = vec![512, 256];

    let dqn = WorkingDQN::new(config)?;
    let device = dqn.device().clone();

    // Test with various portfolio values
    let portfolio_values = vec![50_000.0, 100_000.0, 150_000.0, 200_000.0];

    for &portfolio_val in &portfolio_values {
        let tracker = PortfolioTracker::new(100_000.0, 0.0001, 0.0);

        // Normalized value should be ratio to initial_capital
        let normalized_value = (portfolio_val / 100_000.0) as f32;

        let mut state_vec = vec![0.0f32; 128];
        state_vec[0] = normalized_value; // After Bug #20 fix, this should be normalized
        state_vec[1] = 0.5;
        state_vec[2] = 0.0001;

        let state = Tensor::from_vec(state_vec, (1, 128), &device)?.to_dtype(DType::F32)?;
        let q_values = dqn.forward(&state)?;
        let q_vec = q_values.flatten_all()?.to_vec1::<f32>()?;

        let max_q = q_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let min_q = q_vec.iter().copied().fold(f32::INFINITY, f32::min);

        println!(
            "Portfolio ${:.0}K (normalized {:.2}): Q-range=[{:.2}, {:.2}]",
            portfolio_val / 1000.0,
            normalized_value,
            min_q,
            max_q
        );

        // Q-values should stay reasonable (NOT explode to 1000-4197)
        assert!(
            max_q.abs() < 5000.0,
            "Q-values exploded with portfolio ${}: max_q={}",
            portfolio_val,
            max_q
        );
    }

    println!("✅ Q-values stay stable across different portfolio values");

    Ok(())
}

#[test]
fn test_clamp_removal_allows_large_q_values() -> Result<()> {
    // Bug #19: Test that Q-values are not artificially clamped
    // Bug #20: Test that normalized portfolio prevents explosions
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.num_actions = 45;
    config.hidden_dims = vec![512, 256];

    let dqn = WorkingDQN::new(config)?;
    let tracker = PortfolioTracker::new(100_000.0, 0.0001, 0.0);
    let device = dqn.device().clone();

    // Test with normalized portfolio features
    let mut state_vec = vec![0.0f32; 128];
    let portfolio_features = tracker.get_portfolio_features(4000.0);
    state_vec[0..3].copy_from_slice(&portfolio_features);

    // Add large market features to test clamp removal
    for i in 3..128 {
        state_vec[i] = (i as f32 * 0.1).sin() * 10.0;
    }

    let state = Tensor::from_vec(state_vec, (1, 128), &device)?.to_dtype(DType::F32)?;
    let q_values = dqn.forward(&state)?;
    let q_vec = q_values.flatten_all()?.to_vec1::<f32>()?;

    let max_q = q_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let min_q = q_vec.iter().copied().fold(f32::INFINITY, f32::min);

    println!("Q-value range with normalized portfolio + large features: [{:.2}, {:.2}]", min_q, max_q);

    // After Bug #19 fix: Q-values are not clamped to ±1000
    // After Bug #20 fix: Q-values don't explode due to normalized portfolio
    assert!(q_vec.len() == 45, "Should have 45 Q-values");

    // All Q-values should be finite (no explosion)
    for &q in &q_vec {
        assert!(q.is_finite(), "Q-value should be finite: {}", q);
    }

    println!("✅ Q-values are not clamped and remain stable");

    Ok(())
}

#[test]
fn test_portfolio_normalization_prevents_feature_imbalance() -> Result<()> {
    // Bug #20: Test that portfolio features are normalized
    let tracker = PortfolioTracker::new(100_000.0, 0.0001, 0.0);
    let features = tracker.get_portfolio_features(4000.0);

    // Bug #20 fix: Portfolio value should be normalized to ~1.0
    // NOT raw $100,000 value
    let portfolio_feature = features[0];

    println!("Portfolio feature value: {}", portfolio_feature);

    // After fix, should be ~1.0 (normalized)
    // Before fix, would be 100,000.0 (raw value)
    assert!(
        portfolio_feature.abs() < 10.0,
        "Portfolio feature should be normalized, got: {}",
        portfolio_feature
    );

    // Check feature scale consistency
    let max_feature = features.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let min_feature = features.iter().copied().fold(f32::INFINITY, f32::min);
    let scale_ratio = max_feature / min_feature.abs().max(0.001);

    println!("Feature scale ratio: {:.2}", scale_ratio);

    // Scale ratio should be reasonable (NOT 100,000x)
    assert!(
        scale_ratio < 1000.0,
        "Feature scale ratio too large: {}",
        scale_ratio
    );

    println!("✅ Portfolio features are properly normalized");

    Ok(())
}