//! End-to-End Integration Tests for Continuous PPO
//!
//! Comprehensive integration tests covering:
//! - Full training loop execution
//! - Convergence verification
//! - Action distribution stability
//! - Checkpoint save/load correctness
//! - Continuous vs discrete comparison

mod ppo_continuous_test_helpers;

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

use anyhow::Result;
use ml::ppo::continuous_policy::ContinuousPolicyConfig;
use ml::ppo::continuous_ppo::{ContinuousPPO, ContinuousPPOConfig};
use ml::ppo::gae::GAEConfig;

#[test]
fn test_continuous_ppo_full_training_loop() -> Result<()> {
    println!("\n=== Test: Full Training Loop (5 epochs) ===");

    let config = create_test_config(64);

    // Train for 5 epochs
    let (ppo, metrics) = train_continuous_ppo(5, config, 10, 20)?;

    // Verify no crashes/errors occurred
    assert_eq!(
        metrics.policy_losses.len(),
        5,
        "Should have 5 policy loss values"
    );
    assert_eq!(
        metrics.value_losses.len(),
        5,
        "Should have 5 value loss values"
    );

    // Verify all losses are finite
    for (i, &loss) in metrics.policy_losses.iter().enumerate() {
        assert!(
            loss.is_finite(),
            "Policy loss at epoch {} should be finite",
            i
        );
    }

    for (i, &loss) in metrics.value_losses.iter().enumerate() {
        assert!(
            loss.is_finite(),
            "Value loss at epoch {} should be finite",
            i
        );
    }

    // Verify training steps incremented
    assert_eq!(ppo.get_training_steps(), 5, "Should have 5 training steps");

    println!("✓ Full training loop completed successfully");
    println!("  Final policy loss: {:.4}", metrics.policy_losses.last().unwrap());
    println!("  Final value loss: {:.4}", metrics.value_losses.last().unwrap());
    println!("  Training steps: {}", ppo.get_training_steps());

    Ok(())
}

#[test]
fn test_continuous_ppo_convergence_5_epochs() -> Result<()> {
    println!("\n=== Test: Loss Convergence (5 epochs) ===");

    let config = create_test_config(64);

    // Train for 5 epochs
    let (_ppo, metrics) = train_continuous_ppo(5, config, 10, 20)?;

    // Verify policy loss decreases or stays stable
    let first_policy_loss = metrics.policy_losses[0];
    let last_policy_loss = *metrics.policy_losses.last().unwrap();

    println!("  Policy loss: {:.4} → {:.4}", first_policy_loss, last_policy_loss);

    // Allow for some variation, but should generally decrease or stabilize
    // (In early epochs, loss might fluctuate due to exploration)
    assert!(
        last_policy_loss < first_policy_loss * 1.5,
        "Policy loss should not increase dramatically"
    );

    // Verify value loss decreases or stays stable
    let first_value_loss = metrics.value_losses[0];
    let last_value_loss = *metrics.value_losses.last().unwrap();

    println!("  Value loss: {:.4} → {:.4}", first_value_loss, last_value_loss);

    assert!(
        last_value_loss < first_value_loss * 1.5,
        "Value loss should not increase dramatically"
    );

    // Verify average reward is improving or stable
    let first_reward = metrics.avg_rewards[0];
    let last_reward = *metrics.avg_rewards.last().unwrap();

    println!("  Avg reward: {:.4} → {:.4}", first_reward, last_reward);

    println!("✓ Convergence verification passed");

    Ok(())
}

#[test]
fn test_continuous_ppo_action_distribution_stability() -> Result<()> {
    println!("\n=== Test: Action Distribution Stability (10 epochs) ===");

    let config = create_test_config(64);

    // Train for 10 epochs
    let (ppo, metrics) = train_continuous_ppo(10, config.clone(), 10, 20)?;

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

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

    // Verify actions are within bounds [0, 1]
    for (i, &action) in actions.iter().enumerate() {
        assert!(
            action >= 0.0 && action <= 1.0,
            "Action {} ({:.4}) should be in [0, 1]",
            i,
            action
        );
    }

    // Compute mean and std dev of actions
    let mean = actions.iter().sum::<f32>() / actions.len() as f32;
    let variance = actions.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / actions.len() as f32;
    let std_dev = variance.sqrt();

    println!("  Action distribution after 10 epochs:");
    println!("    Mean: {:.4}", mean);
    println!("    Std Dev: {:.4}", std_dev);
    println!("    Min: {:.4}", actions.iter().cloned().fold(f32::INFINITY, f32::min));
    println!("    Max: {:.4}", actions.iter().cloned().fold(f32::NEG_INFINITY, f32::max));

    // Verify mean is reasonable (not stuck at boundary)
    assert!(
        mean > 0.1 && mean < 0.9,
        "Mean action ({:.4}) should not be stuck at boundary",
        mean
    );

    // Verify some exploration is happening (std dev > 0)
    assert!(std_dev > 0.01, "Std dev ({:.4}) should show exploration", std_dev);

    // Verify exploration decreases over training
    assert!(
        metrics.is_exploration_decaying(),
        "Exploration (log_std) should decrease over training"
    );

    println!("  Exploration decay:");
    println!("    Initial log_std: {:.4}", metrics.log_stds[0]);
    println!("    Final log_std: {:.4}", metrics.log_stds.last().unwrap());

    println!("✓ Action distribution is stable and well-behaved");

    Ok(())
}

#[test]
fn test_continuous_ppo_checkpoint_save_load() -> Result<()> {
    println!("\n=== Test: Checkpoint Save/Load ===");

    let config = create_test_config(64);

    // Train for 5 epochs, save checkpoint
    let (ppo1, metrics1) = train_continuous_ppo(5, config.clone(), 10, 20)?;

    // Get action from trained model
    let test_state = vec![0.5; config.state_dim];
    let (action1, value1) = ppo1.act(&test_state)?;

    println!("  Before save:");
    println!("    Action: {:.4}", action1.position_size());
    println!("    Value: {:.4}", value1);
    println!("    Training steps: {}", ppo1.get_training_steps());

    // Save checkpoints
    let checkpoint_dir = std::path::PathBuf::from("/tmp/foxhunt_continuous_ppo_test");
    std::fs::create_dir_all(&checkpoint_dir)?;

    let actor_path = checkpoint_dir.join("actor.safetensors");
    let critic_path = checkpoint_dir.join("critic.safetensors");

    ppo1.actor.save(&actor_path)?;
    ppo1.critic.save(&critic_path)?;

    println!("  Checkpoints saved to: {}", checkpoint_dir.display());

    // Create new PPO instance
    let ppo2 = ContinuousPPO::new(config.clone())?;

    // Load checkpoints
    ppo2.actor.load(&actor_path)?;
    ppo2.critic.load(&critic_path)?;

    println!("  Checkpoints loaded successfully");

    // Get action from loaded model (should be identical)
    let (action2, value2) = ppo2.act(&test_state)?;

    println!("  After load:");
    println!("    Action: {:.4}", action2.position_size());
    println!("    Value: {:.4}", value2);

    // Verify actions are identical (deterministic with same state)
    assert!(
        (action1.position_size() - action2.position_size()).abs() < 1e-4,
        "Actions should match after loading checkpoint"
    );
    assert!(
        (value1 - value2).abs() < 1e-4,
        "Values should match after loading checkpoint"
    );

    // Continue training from loaded checkpoint
    let (ppo3, metrics3) = {
        let mut ppo = ppo2;
        let mut env = SimpleTradingEnv::new(200, config.state_dim);
        let mut metrics = TrainingMetrics::new();

        // Train for 5 more epochs
        for epoch in 0..5 {
            let mut trajectories = Vec::new();
            let mut epoch_rewards = 0.0;

            for _ in 0..10 {
                let trajectory = ppo_continuous_test_helpers::collect_trajectory(
                    &ppo,
                    &mut env,
                    20,
                )?;
                epoch_rewards += trajectory.steps().iter().map(|s| s.reward).sum::<f32>();
                trajectories.push(trajectory);
            }

            let (advantages, returns) =
                ppo_continuous_test_helpers::compute_batch_gae(&trajectories, &config.gae_config)?;

            let mut batch =
                ml::ppo::continuous_ppo::ContinuousTrajectoryBatch::from_trajectories(
                    trajectories,
                    advantages,
                    returns,
                );

            let (policy_loss, value_loss) = ppo.update(&mut batch)?;

            let test_state = vec![0.0; config.state_dim];
            let log_std = ppo.get_exploration_param(&test_state)?;

            let avg_reward = epoch_rewards / 10.0;
            metrics.add_epoch(policy_loss, value_loss, avg_reward, log_std);

            if (epoch + 1) % 2 == 0 {
                println!("    Continued epoch {}/5: Policy Loss={:.4}, Value Loss={:.4}",
                    epoch + 1, policy_loss, value_loss);
            }
        }

        (ppo, metrics)
    };

    // Verify training continued correctly
    assert_eq!(
        ppo3.get_training_steps(),
        5,
        "Training steps should reset for new instance"
    );

    // Verify losses are still finite
    for loss in &metrics3.policy_losses {
        assert!(loss.is_finite(), "Policy loss should remain finite after resume");
    }

    println!("✓ Checkpoint save/load/resume working correctly");

    // Cleanup
    std::fs::remove_dir_all(&checkpoint_dir)?;

    Ok(())
}

#[test]
fn test_continuous_vs_discrete_comparison_lightweight() -> Result<()> {
    println!("\n=== Test: Continuous vs Discrete Comparison (Lightweight) ===");

    let state_dim = 64;

    // Train continuous PPO
    let continuous_config = create_test_config(state_dim);
    let (continuous_ppo, continuous_metrics) =
        train_continuous_ppo(10, continuous_config, 10, 20)?;

    // Compute simple metrics for continuous
    let continuous_final_reward = *continuous_metrics.avg_rewards.last().unwrap();
    let continuous_policy_loss = *continuous_metrics.policy_losses.last().unwrap();

    println!("\n  Continuous PPO Results:");
    println!("    Final avg reward: {:.4}", continuous_final_reward);
    println!("    Final policy loss: {:.4}", continuous_policy_loss);
    println!("    Exploration decay: {:.4} → {:.4}",
        continuous_metrics.log_stds[0],
        continuous_metrics.log_stds.last().unwrap()
    );

    // For discrete comparison, we'd need to train a discrete PPO model
    // For now, just verify continuous PPO achieves reasonable performance
    println!("\n  Note: Full discrete comparison requires discrete PPO training");
    println!("  Verifying continuous PPO achieves reasonable performance:");

    // Verify continuous PPO learned something
    assert!(
        continuous_metrics.is_policy_loss_converging(0.1) ||
        continuous_metrics.policy_losses.last().unwrap() < &10.0,
        "Continuous PPO should show learning (converging loss or low final loss)"
    );

    // Sample actions from final policy to verify granularity
    let mut actions = Vec::new();
    let test_state = vec![0.5; state_dim];
    for _ in 0..100 {
        let (action, _value) = continuous_ppo.act(&test_state)?;
        actions.push(action.position_size());
    }

    // Count unique action values (quantized to 2 decimals)
    let mut unique_actions = std::collections::HashSet::new();
    for action in &actions {
        let quantized = (action * 100.0).round() as i32;
        unique_actions.insert(quantized);
    }

    println!("  Action granularity: {} unique values (out of 100 samples)", unique_actions.len());

    // Verify continuous action space provides granularity
    assert!(
        unique_actions.len() >= 5,
        "Continuous actions should show granularity (found {} unique values)",
        unique_actions.len()
    );

    println!("✓ Continuous PPO provides fine-grained action control");

    Ok(())
}

#[test]
fn test_continuous_ppo_with_real_data_subset() -> Result<()> {
    println!("\n=== Test: Training with Real Data Subset ===");

    // Load 500 samples from ES_FUT data
    let prices = ppo_continuous_test_helpers::load_test_data(500)?;
    println!("  Loaded {} price samples from ES_FUT_180d.parquet", prices.len());

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

    // Create environment with real prices
    let mut env = SimpleTradingEnv::from_prices(prices, state_dim);

    // Create PPO
    let mut ppo = ContinuousPPO::new(config.clone())?;

    // Train for 10 epochs
    let mut metrics = TrainingMetrics::new();

    for epoch in 0..10 {
        let mut trajectories = Vec::new();
        let mut epoch_rewards = 0.0;

        for _ in 0..10 {
            let trajectory = ppo_continuous_test_helpers::collect_trajectory(&ppo, &mut env, 50)?;
            epoch_rewards += trajectory.steps().iter().map(|s| s.reward).sum::<f32>();
            trajectories.push(trajectory);
        }

        let (advantages, returns) =
            ppo_continuous_test_helpers::compute_batch_gae(&trajectories, &config.gae_config)?;

        let mut batch = ml::ppo::continuous_ppo::ContinuousTrajectoryBatch::from_trajectories(
            trajectories,
            advantages,
            returns,
        );

        let (policy_loss, value_loss) = ppo.update(&mut batch)?;

        let test_state = vec![0.0; state_dim];
        let log_std = ppo.get_exploration_param(&test_state)?;

        let avg_reward = epoch_rewards / 10.0;
        metrics.add_epoch(policy_loss, value_loss, avg_reward, log_std);

        if (epoch + 1) % 3 == 0 {
            println!("  Epoch {}/10: Policy Loss={:.4}, Value Loss={:.4}, Avg Reward={:.4}",
                epoch + 1, policy_loss, value_loss, avg_reward);
        }
    }

    // Verify training completed successfully
    assert_eq!(metrics.policy_losses.len(), 10, "Should have 10 epochs");

    println!("\n  Training Summary:");
    println!("    Initial policy loss: {:.4}", metrics.policy_losses[0]);
    println!("    Final policy loss: {:.4}", metrics.policy_losses.last().unwrap());
    println!("    Initial avg reward: {:.4}", metrics.avg_rewards[0]);
    println!("    Final avg reward: {:.4}", metrics.avg_rewards.last().unwrap());

    println!("✓ Training with real data completed successfully");

    Ok(())
}