// Basic test to verify LSTM-PPO can complete a training update
use ml::ppo::ppo::{PPOConfig, PPO};
use ml::ppo::trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep};
use ml::dqn::TradingAction;
use candle_core::Device;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Testing LSTM-PPO training loop...");
    
    // Create LSTM-enabled PPO config
    let mut config = PPOConfig::default();
    config.use_lstm = true;
    config.lstm_hidden_dim = 64;
    config.lstm_num_layers = 1;
    config.lstm_sequence_length = 8;
    config.state_dim = 10;
    config.num_actions = 3;
    config.num_epochs = 1;
    config.batch_size = 32;
    config.mini_batch_size = 16;
    
    let device = Device::Cpu;
    
    // Create LSTM-PPO agent
    println!("Creating LSTM-PPO agent...");
    let mut ppo = PPO::with_device(config.clone(), device.clone())?;
    
    // Create a trajectory
    println!("Creating trajectory...");
    let mut trajectory = Trajectory::new();
    
    // Add steps to trajectory
    for i in 0..32 {
        let state = vec![0.1; 10];
        let action = TradingAction::Hold;
        let log_prob = -1.0;
        let value = 0.5;
        let reward = 1.0;
        let done = i == 31; // Last step is done
        
        let step = TrajectoryStep::new(state, action, log_prob, value, reward, done);
        trajectory.add_step(step);
    }
    
    // Compute advantages using GAE
    let mut advantages = vec![0.0; 32];
    let mut returns = vec![0.0; 32];
    
    let gamma = config.gae_config.gamma;
    let lambda = config.gae_config.lambda;
    
    // Simple GAE computation
    let mut next_value = 0.0;
    let mut next_advantage = 0.0;
    for i in (0..32).rev() {
        let reward = 1.0;
        let value = 0.5;
        let done = i == 31;
        
        let delta = reward + gamma * next_value * (1.0 - done as i32 as f32) - value;
        advantages[i] = delta + gamma * lambda * next_advantage * (1.0 - done as i32 as f32);
        returns[i] = advantages[i] + value;
        
        next_value = value;
        next_advantage = advantages[i];
    }
    
    // Create batch from trajectory
    let mut batch = TrajectoryBatch::from_trajectories(
        vec![trajectory],
        advantages,
        returns,
    );
    
    // Normalize advantages
    println!("Normalizing advantages...");
    batch.normalize_advantages()?;
    
    // Run LSTM training update
    println!("Running LSTM training update...");
    let (policy_loss, value_loss) = ppo.update(&mut batch)?;
    
    println!("LSTM-PPO training completed successfully!");
    println!("  Policy loss: {}", policy_loss);
    println!("  Value loss: {}", value_loss);
    println!("  Losses are finite: {}", policy_loss.is_finite() && value_loss.is_finite());
    
    if !policy_loss.is_finite() || !value_loss.is_finite() {
        return Err("NaN/Inf detected in losses".into());
    }
    
    println!("\nTest PASSED: LSTM-PPO can complete a training update");
    Ok(())
}