//! Performance Benchmark Tests for Continuous PPO
//!
//! Benchmarks covering:
//! - Sharpe ratio computation
//! - Training time measurements
//! - GPU vs CPU comparison
//! - Throughput metrics

mod ppo_continuous_test_helpers;

use ppo_continuous_test_helpers::{
    create_test_config, load_test_data, train_continuous_ppo, SimpleTradingEnv,
};

use anyhow::Result;
use ml::ppo::continuous_ppo::ContinuousPPO;
use std::time::Instant;

/// Compute Sharpe ratio from returns
fn compute_sharpe_ratio(returns: &[f32]) -> f32 {
    if returns.is_empty() {
        return 0.0;
    }

    let mean = returns.iter().sum::<f32>() / returns.len() as f32;
    let variance = returns
        .iter()
        .map(|&r| (r - mean).powi(2))
        .sum::<f32>()
        / returns.len() as f32;
    let std = variance.sqrt();

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

    // Annualized Sharpe (assuming daily returns, 252 trading days)
    mean / std * (252.0_f32).sqrt()
}

/// Compute maximum drawdown from equity curve
fn compute_max_drawdown(equity: &[f32]) -> f32 {
    if equity.is_empty() {
        return 0.0;
    }

    let mut peak = equity[0];
    let mut max_dd = 0.0;

    for &value in equity {
        if value > peak {
            peak = value;
        }
        let drawdown = (peak - value) / peak.max(1e-8);
        max_dd = max_dd.max(drawdown);
    }

    max_dd
}

/// Run backtest and compute trading metrics
fn run_backtest(ppo: &ContinuousPPO, prices: &[f32], state_dim: usize) -> Result<BacktestMetrics> {
    let mut env = SimpleTradingEnv::from_prices(prices.to_vec(), state_dim);
    let mut returns = Vec::new();
    let mut equity = vec![10000.0]; // Start with $10,000
    let mut wins = 0;
    let mut losses = 0;

    let state = env.reset();
    let mut current_state = state;

    while !env.is_done() {
        // Get action from policy
        let (action, _value) = ppo.act(&current_state)?;

        // Execute action
        let (next_state, reward, done) = env.step(&action);

        // Track equity
        let current_equity = *equity.last().unwrap();
        let new_equity = current_equity + reward;
        equity.push(new_equity);

        // Track returns
        let ret = reward / current_equity;
        returns.push(ret);

        // Track wins/losses
        if reward > 0.0 {
            wins += 1;
        } else if reward < 0.0 {
            losses += 1;
        }

        current_state = next_state;

        if done {
            break;
        }
    }

    let sharpe = compute_sharpe_ratio(&returns);
    let max_dd = compute_max_drawdown(&equity);
    let total_return = (*equity.last().unwrap() - equity[0]) / equity[0];
    let win_rate = if wins + losses > 0 {
        wins as f32 / (wins + losses) as f32
    } else {
        0.0
    };

    Ok(BacktestMetrics {
        sharpe_ratio: sharpe,
        total_return,
        max_drawdown: max_dd,
        win_rate,
        num_trades: wins + losses,
        final_equity: *equity.last().unwrap(),
    })
}

#[derive(Debug, Clone)]
struct BacktestMetrics {
    sharpe_ratio: f32,
    total_return: f32,
    max_drawdown: f32,
    win_rate: f32,
    num_trades: usize,
    final_equity: f32,
}

#[test]
fn test_continuous_ppo_sharpe_ratio() -> Result<()> {
    println!("\n=== Test: Sharpe Ratio Benchmark ===");

    // Load real data
    let prices = load_test_data(1000)?;
    println!("  Loaded {} price samples", prices.len());

    let state_dim = 64;
    let config = create_test_config(state_dim);

    // Train continuous PPO
    println!("  Training continuous PPO for 50 epochs...");
    let (ppo, metrics) = train_continuous_ppo(50, config.clone(), 20, 50)?;

    println!("\n  Training Results:");
    println!("    Final policy loss: {:.4}", metrics.policy_losses.last().unwrap());
    println!("    Final value loss: {:.4}", metrics.value_losses.last().unwrap());
    println!("    Final avg reward: {:.4}", metrics.avg_rewards.last().unwrap());

    // Run backtest
    println!("\n  Running backtest on test data...");
    let backtest_metrics = run_backtest(&ppo, &prices, state_dim)?;

    println!("\n  Backtest Results:");
    println!("    Sharpe Ratio: {:.4}", backtest_metrics.sharpe_ratio);
    println!("    Total Return: {:.2}%", backtest_metrics.total_return * 100.0);
    println!("    Max Drawdown: {:.2}%", backtest_metrics.max_drawdown * 100.0);
    println!("    Win Rate: {:.2}%", backtest_metrics.win_rate * 100.0);
    println!("    Num Trades: {}", backtest_metrics.num_trades);
    println!("    Final Equity: ${:.2}", backtest_metrics.final_equity);

    // Note: Discrete baseline is 4.311 from CLAUDE.md (Wave 7 best DQN result)
    // For continuous PPO, we expect competitive but different performance
    // due to different action space (continuous position sizing vs discrete actions)

    // Verify Sharpe ratio is reasonable (not NaN, finite)
    assert!(
        backtest_metrics.sharpe_ratio.is_finite(),
        "Sharpe ratio should be finite"
    );

    // Verify win rate is reasonable (>= 30%)
    assert!(
        backtest_metrics.win_rate >= 0.3,
        "Win rate should be >= 30% (got {:.2}%)",
        backtest_metrics.win_rate * 100.0
    );

    // Verify max drawdown is reasonable (< 50%)
    assert!(
        backtest_metrics.max_drawdown < 0.5,
        "Max drawdown should be < 50% (got {:.2}%)",
        backtest_metrics.max_drawdown * 100.0
    );

    println!("\n  Performance Assessment:");
    if backtest_metrics.sharpe_ratio >= 3.0 {
        println!("    ✓ Excellent: Sharpe ≥ 3.0 (comparable to discrete baseline 4.311)");
    } else if backtest_metrics.sharpe_ratio >= 1.5 {
        println!("    ✓ Good: Sharpe ≥ 1.5 (reasonable for continuous action space)");
    } else if backtest_metrics.sharpe_ratio >= 0.5 {
        println!("    ⚠ Acceptable: Sharpe ≥ 0.5 (learning occurred, room for improvement)");
    } else {
        println!("    ⚠ Needs improvement: Sharpe < 0.5");
    }

    println!("\n✓ Sharpe ratio benchmark complete");

    Ok(())
}

#[test]
fn test_continuous_ppo_training_time() -> Result<()> {
    println!("\n=== Test: Training Time Benchmark ===");

    let config = create_test_config(64);

    // Benchmark training time per epoch
    println!("  Benchmarking 10 epochs...");
    let start = Instant::now();
    let (_ppo, _metrics) = train_continuous_ppo(10, config, 10, 20)?;
    let elapsed = start.elapsed();

    let time_per_epoch = elapsed.as_secs_f32() / 10.0;

    println!("\n  Timing Results:");
    println!("    Total time: {:.2}s", elapsed.as_secs_f32());
    println!("    Time per epoch: {:.3}s", time_per_epoch);
    println!("    Throughput: {:.1} epochs/min", 60.0 / time_per_epoch);

    // Verify training is reasonably fast
    // Note: Discrete DQN baseline is ~15s total (from CLAUDE.md)
    // Continuous PPO is expected to be slower due to:
    // - More complex policy network (mean + log_std outputs)
    // - Gaussian sampling
    // - GAE computation
    // Target: < 2x discrete DQN time

    assert!(
        time_per_epoch < 5.0,
        "Training should be < 5s per epoch (got {:.3}s)",
        time_per_epoch
    );

    println!("\n  Performance Assessment:");
    if time_per_epoch < 0.5 {
        println!("    ✓ Excellent: < 0.5s per epoch");
    } else if time_per_epoch < 1.0 {
        println!("    ✓ Good: < 1.0s per epoch");
    } else if time_per_epoch < 2.0 {
        println!("    ✓ Acceptable: < 2.0s per epoch");
    } else {
        println!("    ⚠ Slow: ≥ 2.0s per epoch (acceptable for more granular control)");
    }

    println!("\n✓ Training time benchmark complete");

    Ok(())
}

#[test]
fn test_continuous_ppo_inference_latency() -> Result<()> {
    println!("\n=== Test: Inference Latency Benchmark ===");

    let state_dim = 64;
    let config = create_test_config(state_dim);

    // Create and train PPO
    let (ppo, _metrics) = train_continuous_ppo(10, config.clone(), 10, 20)?;

    // Warm-up
    let test_state = vec![0.5; state_dim];
    for _ in 0..100 {
        let _ = ppo.act(&test_state)?;
    }

    // Benchmark inference
    let num_inferences = 10000;
    let start = Instant::now();

    for _ in 0..num_inferences {
        let _ = ppo.act(&test_state)?;
    }

    let elapsed = start.elapsed();
    let latency_us = elapsed.as_micros() as f32 / num_inferences as f32;

    println!("\n  Inference Latency:");
    println!("    Total inferences: {}", num_inferences);
    println!("    Total time: {:.3}s", elapsed.as_secs_f32());
    println!("    Average latency: {:.1}μs", latency_us);
    println!("    Throughput: {:.1} inferences/sec", num_inferences as f32 / elapsed.as_secs_f32());

    // Note: Discrete PPO baseline is ~324μs (from CLAUDE.md)
    // Continuous PPO should be comparable or slightly slower

    // Verify inference is fast enough for trading
    // Target: < 1ms (1000μs) for HFT compatibility
    assert!(
        latency_us < 1000.0,
        "Inference should be < 1ms (got {:.1}μs)",
        latency_us
    );

    println!("\n  Performance Assessment:");
    if latency_us < 100.0 {
        println!("    ✓ Excellent: < 100μs (sub-millisecond latency)");
    } else if latency_us < 300.0 {
        println!("    ✓ Good: < 300μs (comparable to discrete baseline 324μs)");
    } else if latency_us < 500.0 {
        println!("    ✓ Acceptable: < 500μs (HFT compatible)");
    } else {
        println!("    ⚠ Slow: ≥ 500μs (still < 1ms target)");
    }

    println!("\n✓ Inference latency benchmark complete");

    Ok(())
}

#[test]
fn test_continuous_ppo_memory_usage() -> Result<()> {
    println!("\n=== Test: Memory Usage Benchmark ===");

    let state_dim = 64;
    let config = create_test_config(state_dim);

    // Create PPO
    let (ppo, _metrics) = train_continuous_ppo(10, config.clone(), 10, 20)?;

    // Estimate parameter count
    // Policy network: state_dim -> 64 -> 32 -> 2 (mean + log_std)
    let policy_params = (state_dim * 64) + 64 +  // Layer 1
                        (64 * 32) + 32 +          // Layer 2
                        (32 * 1) + 1 +            // Mean head
                        (32 * 1) + 1;             // Log std head

    // Value network: state_dim -> 64 -> 32 -> 1
    let value_params = (state_dim * 64) + 64 +   // Layer 1
                       (64 * 32) + 32 +           // Layer 2
                       (32 * 1) + 1;              // Output head

    let total_params = policy_params + value_params;

    // Estimate memory (FP32 = 4 bytes per param)
    let memory_mb = (total_params * 4) as f32 / 1_000_000.0;

    println!("\n  Model Size:");
    println!("    Policy network params: {}", policy_params);
    println!("    Value network params: {}", value_params);
    println!("    Total params: {}", total_params);
    println!("    Estimated memory (FP32): {:.2} MB", memory_mb);

    // Note: Discrete DQN baseline is ~6MB, PPO is ~145MB (from CLAUDE.md)
    // Continuous PPO should be similar to discrete PPO

    // Verify memory usage is reasonable (< 200MB)
    assert!(
        memory_mb < 200.0,
        "Memory usage should be < 200MB (got {:.2}MB)",
        memory_mb
    );

    // Verify model loaded successfully
    let test_state = vec![0.5; state_dim];
    let (action, value) = ppo.act(&test_state)?;

    println!("\n  Model Verification:");
    println!("    Sample action: {:.4}", action.position_size());
    println!("    Sample value: {:.4}", value);

    assert!(
        action.position_size() >= 0.0 && action.position_size() <= 1.0,
        "Action should be in valid range"
    );
    assert!(value.is_finite(), "Value should be finite");

    println!("\n  Memory Assessment:");
    if memory_mb < 50.0 {
        println!("    ✓ Excellent: < 50MB (very lightweight)");
    } else if memory_mb < 100.0 {
        println!("    ✓ Good: < 100MB (lightweight)");
    } else if memory_mb < 150.0 {
        println!("    ✓ Acceptable: < 150MB (reasonable for continuous PPO)");
    } else {
        println!("    ⚠ Large: ≥ 150MB (still < 200MB target)");
    }

    println!("\n✓ Memory usage benchmark complete");

    Ok(())
}

#[test]
fn test_continuous_vs_discrete_action_granularity() -> Result<()> {
    println!("\n=== Test: Action Granularity Comparison ===");

    let state_dim = 64;
    let config = create_test_config(state_dim);

    // Train continuous PPO
    let (ppo, _metrics) = train_continuous_ppo(20, config.clone(), 10, 20)?;

    // Sample 1000 actions
    let test_state = vec![0.5; state_dim];
    let mut actions = Vec::new();

    for _ in 0..1000 {
        let (action, _value) = ppo.act(&test_state)?;
        actions.push(action.position_size());
    }

    // Count unique actions (quantized to 3 decimals = 1000 levels)
    let mut unique_actions_fine = std::collections::HashSet::new();
    for action in &actions {
        let quantized = (action * 1000.0).round() as i32;
        unique_actions_fine.insert(quantized);
    }

    // Count unique actions (quantized to 1 decimal = 10 levels, like discrete)
    let mut unique_actions_coarse = std::collections::HashSet::new();
    for action in &actions {
        let quantized = (action * 10.0).round() as i32;
        unique_actions_coarse.insert(quantized);
    }

    println!("\n  Action Granularity:");
    println!("    Samples: 1000");
    println!("    Unique actions (fine, 0.001 resolution): {}", unique_actions_fine.len());
    println!("    Unique actions (coarse, 0.1 resolution): {}", unique_actions_coarse.len());
    println!("    Mean: {:.4}", actions.iter().sum::<f32>() / actions.len() as f32);
    println!("    Std: {:.4}", {
        let mean = actions.iter().sum::<f32>() / actions.len() as f32;
        let var = actions.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / actions.len() as f32;
        var.sqrt()
    });

    // Verify continuous actions provide more granularity than discrete
    assert!(
        unique_actions_fine.len() >= 20,
        "Continuous actions should provide fine granularity (got {} unique values)",
        unique_actions_fine.len()
    );

    assert!(
        unique_actions_coarse.len() >= 5,
        "Continuous actions should span multiple coarse levels (got {} levels)",
        unique_actions_coarse.len()
    );

    println!("\n  Comparison to Discrete (45-action space):");
    println!("    Discrete: 45 possible actions");
    println!("    Continuous (fine): {} effective levels", unique_actions_fine.len());
    println!("    Granularity advantage: {:.1}x",
        unique_actions_fine.len() as f32 / 45.0
    );

    println!("\n✓ Continuous actions provide fine-grained control");

    Ok(())
}

#[test]
fn test_scalability_with_batch_size() -> Result<()> {
    println!("\n=== Test: Scalability with Batch Size ===");

    let state_dim = 64;

    // Test different batch sizes
    let batch_sizes = vec![16, 32, 64, 128];

    for &batch_size in &batch_sizes {
        let mut config = create_test_config(state_dim);
        config.batch_size = batch_size;
        config.mini_batch_size = (batch_size / 2).max(8);

        println!("\n  Testing batch_size={}, mini_batch_size={}",
            batch_size, config.mini_batch_size);

        let start = Instant::now();
        let (_ppo, metrics) = train_continuous_ppo(5, config, 10, 20)?;
        let elapsed = start.elapsed();

        println!("    Time: {:.2}s ({:.3}s/epoch)",
            elapsed.as_secs_f32(),
            elapsed.as_secs_f32() / 5.0
        );
        println!("    Final policy loss: {:.4}", metrics.policy_losses.last().unwrap());
        println!("    Final value loss: {:.4}", metrics.value_losses.last().unwrap());

        // Verify training completed successfully
        assert_eq!(metrics.policy_losses.len(), 5, "Should complete 5 epochs");
        for loss in &metrics.policy_losses {
            assert!(loss.is_finite(), "Loss should be finite");
        }
    }

    println!("\n✓ All batch sizes scale successfully");

    Ok(())
}