//! Performance smoke tests.
//!
//! Speed benchmarks for the DQN training pipeline. Requires GPU and
//! real DBN test data in `test_data/ES.FUT/`.
//!
//! ```sh
//! cargo test -p ml --lib smoke_tests::performance -- --nocapture
//! ```
use super::helpers::{assert_finite, smoke_params, test_data_dir};
use crate::trainers::dqn::DQNTrainer;
use ml_core::device::MlDevice;
use std::sync::Arc;
use std::time::Instant;
use tracing::info;

/// Resolve test data dir — hard error if missing.
fn data_dir() -> String {
    test_data_dir()
        .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback")
}

/// Measure end-to-end training throughput (3 epochs, real data).
#[test]
#[ignore] // Loads real training data — run via nightly CI or manual trigger
fn test_training_throughput_measurement() -> anyhow::Result<()> {
    let mut trainer = DQNTrainer::new(smoke_params())?; // auto-detect GPU
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;

    let start = Instant::now();
    let metrics = rt.block_on(trainer.train(&data_dir(), |_epoch, _bytes, _best| {
        Ok("skip".to_owned())
    }))?;
    let elapsed = start.elapsed();

    assert_finite(metrics.loss, "loss");
    info!(
        elapsed_secs = elapsed.as_secs_f64(),
        epochs_trained = metrics.epochs_trained,
        "Training throughput"
    );
    info!(loss = metrics.loss, "Avg loss");
    if let Some(&avg_q) = metrics.additional_metrics.get("avg_q_value") {
        info!(avg_q_value = avg_q, "Avg Q-value");
    }
    if let Some(&final_eps) = metrics.additional_metrics.get("final_epsilon") {
        info!(final_epsilon = final_eps, "Final epsilon");
    }
    drop(trainer);
    drop(rt);
    Ok(())
}

/// Measure per-sample latency of the GPU PER replay buffer.
#[test]
fn test_per_sample_latency() -> anyhow::Result<()> {
    use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};

    let dev = MlDevice::new_cuda(0)?;
    let stream = Arc::clone(dev.cuda_stream()?);
    let state_dim: usize = 48;
    let capacity: usize = 10_000;
    let config = GpuReplayBufferConfig {
        capacity,
        state_dim,
        alpha: 0.6,
        beta_start: 0.4,
        beta_max: 1.0,
        beta_annealing_steps: 1000,
        epsilon: 1e-6,
        max_memory_bytes: 4 * 1024 * 1024 * 1024,
        max_batch_size: 1024,
    };
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    // Fill with 5000 experiences
    let fill_count: usize = 5000;
    let batch_insert: usize = 100;
    for _ in 0..(fill_count / batch_insert) {
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let host_states: Vec<f32> = (0..batch_insert * state_dim).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();
        let host_rewards: Vec<f32> = (0..batch_insert).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();

        let states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?;
        let next_states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?;
        let actions = stream.alloc_zeros::<u32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
        let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?;
        let dones = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
        buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, batch_insert)?;
    }
    assert_eq!(buf.len(), fill_count);

    // Warmup
    for _ in 0..5 {
        let _ = buf.sample_proportional(64)?;
    }

    // Timed: 100 proportional samples of batch_size=64
    let sample_rounds: usize = 100;
    let start = Instant::now();
    for _ in 0..sample_rounds {
        let _ = buf.sample_proportional(64)?;
    }
    let elapsed = start.elapsed();
    let us_per_sample = elapsed.as_micros() as f64 / sample_rounds as f64;

    info!(
        us_per_sample = us_per_sample,
        rounds = sample_rounds,
        batch_size = 64,
        "PER proportional sample latency"
    );
    assert!(
        us_per_sample < 50_000.0,
        "sample latency {us_per_sample:.1} us is unreasonably high"
    );

    Ok(())
}

/// Single-epoch training on real data.
/// Loads full 163K-bar dataset — too slow for CI (run via nightly or manual trigger).
#[test]
#[ignore]
fn test_real_data_single_epoch() -> anyhow::Result<()> {
    let mut params = smoke_params();
    params.epochs = 1;
    let mut trainer = DQNTrainer::new(params)?;
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;

    let start = Instant::now();
    let metrics = rt.block_on(trainer.train(&data_dir(), |_epoch, _bytes, _best| {
        Ok("skip".to_owned())
    }))?;
    let elapsed = start.elapsed();

    assert_finite(metrics.loss, "real_data_loss");
    info!(
        elapsed_secs = elapsed.as_secs_f64(),
        loss = metrics.loss,
        "Real-data single epoch"
    );
    if let Some(&avg_q) = metrics.additional_metrics.get("avg_q_value") {
        info!(avg_q_value = avg_q, "Avg Q-value");
    }
    drop(trainer);
    drop(rt);
    Ok(())
}
