foxhunt/ml/tests/ppo_continuous_test_helpers.rs

//! Test Helpers for Continuous PPO
//!
//! Provides utilities for testing continuous PPO implementation:
//! - Simple trading environment for testing
//! - Data loading utilities
//! - Training metric collection
//! - Trajectory generation helpers

use candle_core::Device;
use ml::ppo::continuous_policy::{ContinuousAction, ContinuousPolicyConfig};
use ml::ppo::continuous_ppo::{
    ContinuousPPO, ContinuousPPOConfig, ContinuousTrajectory, ContinuousTrajectoryBatch,
    ContinuousTrajectoryStep,
};
use ml::ppo::gae::GAEConfig;
use std::fs::File;
use std::path::PathBuf;

/// Training metrics collected during continuous PPO training
#[derive(Debug, Clone)]
pub struct TrainingMetrics {
    /// Policy loss per epoch
    pub policy_losses: Vec<f32>,
    /// Value loss per epoch
    pub value_losses: Vec<f32>,
    /// Average reward per epoch
    pub avg_rewards: Vec<f32>,
    /// Log standard deviation per epoch (exploration parameter)
    pub log_stds: Vec<f32>,
}

impl TrainingMetrics {
    pub fn new() -> Self {
        Self {
            policy_losses: Vec::new(),
            value_losses: Vec::new(),
            avg_rewards: Vec::new(),
            log_stds: Vec::new(),
        }
    }

    /// Add epoch metrics
    pub fn add_epoch(&mut self, policy_loss: f32, value_loss: f32, avg_reward: f32, log_std: f32) {
        self.policy_losses.push(policy_loss);
        self.value_losses.push(value_loss);
        self.avg_rewards.push(avg_reward);
        self.log_stds.push(log_std);
    }

    /// Check if policy loss is converging (decreasing trend)
    pub fn is_policy_loss_converging(&self, min_reduction_pct: f32) -> bool {
        if self.policy_losses.len() < 2 {
            return false;
        }
        let first = self.policy_losses[0];
        let last = *self.policy_losses.last().unwrap();
        let reduction = (first - last) / first.abs().max(1e-8);
        reduction >= min_reduction_pct
    }

    /// Check if value loss is stable
    pub fn is_value_loss_stable(&self, window: usize, max_std: f32) -> bool {
        if self.value_losses.len() < window {
            return false;
        }
        let recent: Vec<f32> = self
            .value_losses
            .iter()
            .rev()
            .take(window)
            .copied()
            .collect();
        let mean = recent.iter().sum::<f32>() / recent.len() as f32;
        let variance = recent.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / recent.len() as f32;
        let std = variance.sqrt();
        std < max_std
    }

    /// Check if rewards are improving
    pub fn are_rewards_improving(&self, min_improvement_factor: f32) -> bool {
        if self.avg_rewards.len() < 2 {
            return false;
        }
        let first = self.avg_rewards[0];
        let last = *self.avg_rewards.last().unwrap();
        last >= first * min_improvement_factor
    }

    /// Check if exploration is decaying (log_std decreasing)
    pub fn is_exploration_decaying(&self) -> bool {
        if self.log_stds.len() < 2 {
            return false;
        }
        let first = self.log_stds[0];
        let last = *self.log_stds.last().unwrap();
        last < first
    }
}

/// Simple trading environment for testing continuous PPO
#[derive(Debug, Clone)]
pub struct SimpleTradingEnv {
    /// Price history
    prices: Vec<f32>,
    /// Current position (continuous value in [0, 1])
    position: f32,
    /// Current step
    step: usize,
    /// State dimension
    state_dim: usize,
    /// Transaction cost (basis points)
    transaction_cost_bps: f32,
}

impl SimpleTradingEnv {
    /// Create new environment with synthetic price data
    pub fn new(num_steps: usize, state_dim: usize) -> Self {
        // Generate simple sine wave prices for testing
        let prices: Vec<f32> = (0..num_steps)
            .map(|i| 100.0 + 10.0 * ((i as f32) * 0.1).sin())
            .collect();

        Self {
            prices,
            position: 0.0,
            step: 0,
            state_dim,
            transaction_cost_bps: 10.0, // 0.10%
        }
    }

    /// Create from real price data
    pub fn from_prices(prices: Vec<f32>, state_dim: usize) -> Self {
        Self {
            prices,
            position: 0.0,
            step: 0,
            state_dim,
            transaction_cost_bps: 10.0,
        }
    }

    /// Reset environment to initial state
    pub fn reset(&mut self) -> Vec<f32> {
        self.step = 0;
        self.position = 0.0;
        self.get_state()
    }

    /// Get current state (simplified features)
    fn get_state(&self) -> Vec<f32> {
        let mut state = vec![0.0; self.state_dim];

        if self.step < self.prices.len() {
            let current_price = self.prices[self.step];

            // Feature 0: Normalized current price
            state[0] = (current_price - 100.0) / 10.0;

            // Feature 1: Price change (if not first step)
            if self.step > 0 {
                let prev_price = self.prices[self.step - 1];
                state[1] = (current_price - prev_price) / prev_price;
            }

            // Feature 2: Current position
            state[2] = self.position;

            // Feature 3: Simple momentum (last 5 steps)
            if self.step >= 5 {
                let momentum = (self.prices[self.step] - self.prices[self.step - 5])
                    / self.prices[self.step - 5];
                state[3] = momentum;
            }

            // Fill remaining features with noise for dimensionality
            for i in 4..self.state_dim {
                state[i] = ((i + self.step) as f32).sin() * 0.1;
            }
        }

        state
    }

    /// Execute action and get reward
    pub fn step(&mut self, action: &ContinuousAction) -> (Vec<f32>, f32, bool) {
        if self.step >= self.prices.len() - 1 {
            return (self.get_state(), 0.0, true);
        }

        let current_price = self.prices[self.step];
        let next_price = self.prices[self.step + 1];

        // Calculate position change
        let new_position = action.position_size();
        let position_change = (new_position - self.position).abs();

        // Calculate transaction costs
        let transaction_cost = position_change * current_price * (self.transaction_cost_bps / 10000.0);

        // Calculate P&L from price movement
        let price_return = (next_price - current_price) / current_price;
        let pnl = self.position * price_return * 100.0; // Scale for easier learning

        // Reward = P&L - transaction costs
        let reward = pnl - transaction_cost;

        // Update state
        self.position = new_position;
        self.step += 1;

        let next_state = self.get_state();
        let done = self.step >= self.prices.len() - 1;

        (next_state, reward, done)
    }

    /// Get current step
    pub fn current_step(&self) -> usize {
        self.step
    }

    /// Check if episode is done
    pub fn is_done(&self) -> bool {
        self.step >= self.prices.len() - 1
    }
}

/// Load subset of ES_FUT parquet data for testing
pub fn load_test_data(num_samples: usize) -> anyhow::Result<Vec<f32>> {
    use parquet::file::reader::{FileReader, SerializedFileReader};

    let parquet_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .join("test_data")
        .join("ES_FUT_180d.parquet");

    let file = File::open(&parquet_path)?;
    let reader = SerializedFileReader::new(file)?;

    let mut prices = Vec::new();

    for row_result in reader.get_row_iter(None)? {
        if prices.len() >= num_samples {
            break;
        }

        let row = row_result?;

        // Extract close price (column 4)
        if let Some(close_field) = row.get_double(4) {
            prices.push(close_field as f32);
        }
    }

    Ok(prices)
}

/// Train continuous PPO for N epochs and return metrics
pub fn train_continuous_ppo(
    epochs: usize,
    config: ContinuousPPOConfig,
    num_trajectories: usize,
    trajectory_length: usize,
) -> anyhow::Result<(ContinuousPPO, TrainingMetrics)> {
    let mut ppo = ContinuousPPO::new(config.clone())?;
    let mut metrics = TrainingMetrics::new();

    // Create environment
    let mut env = SimpleTradingEnv::new(trajectory_length * 10, config.state_dim);

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

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

        // Compute advantages using GAE
        let (advantages, returns) = compute_batch_gae(&trajectories, &config.gae_config)?;

        // Create batch
        let mut batch = ContinuousTrajectoryBatch::from_trajectories(trajectories, advantages, returns);

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

        // Get current exploration parameter (log_std)
        let test_state = vec![0.0; config.state_dim];
        let log_std = ppo.get_exploration_param(&test_state)?;

        // Record metrics
        let avg_reward = epoch_rewards / num_trajectories as f32;
        metrics.add_epoch(policy_loss, value_loss, avg_reward, log_std);

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

    Ok((ppo, metrics))
}

/// Collect a single trajectory
pub fn collect_trajectory(
    ppo: &ContinuousPPO,
    env: &mut SimpleTradingEnv,
    max_steps: usize,
) -> anyhow::Result<ContinuousTrajectory> {
    let mut trajectory = ContinuousTrajectory::new();
    let mut state = env.reset();

    for _ in 0..max_steps {
        if env.is_done() {
            break;
        }

        // Get action with log prob
        let (action, log_prob, value) = ppo.act_with_log_prob(&state)?;

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

        // Add step to trajectory
        trajectory.add_step(ContinuousTrajectoryStep::new(
            state.clone(),
            action,
            log_prob,
            reward,
            value,
            done,
        ));

        state = next_state;

        if done {
            break;
        }
    }

    Ok(trajectory)
}

/// Compute GAE advantages for batch of trajectories
pub fn compute_batch_gae(
    trajectories: &[ContinuousTrajectory],
    gae_config: &GAEConfig,
) -> anyhow::Result<(Vec<f32>, Vec<f32>)> {
    let mut all_advantages = Vec::new();
    let mut all_returns = Vec::new();

    for trajectory in trajectories {
        let steps = trajectory.steps();
        if steps.is_empty() {
            continue;
        }

        let rewards: Vec<f32> = steps.iter().map(|s| s.reward).collect();
        let values: Vec<f32> = steps.iter().map(|s| s.value).collect();
        let dones: Vec<bool> = steps.iter().map(|s| s.done).collect();

        // Last value is 0 if done, otherwise use last step's value
        let next_value = if steps.last().unwrap().done {
            0.0
        } else {
            steps.last().unwrap().value
        };

        // Compute GAE
        let (advantages, returns) =
            ml::ppo::gae::compute_gae_single_trajectory(&rewards, &values, &dones, next_value, gae_config)?;

        all_advantages.extend(advantages);
        all_returns.extend(returns);
    }

    Ok((all_advantages, all_returns))
}

/// Create default continuous PPO config for testing
pub fn create_test_config(state_dim: usize) -> ContinuousPPOConfig {
    ContinuousPPOConfig {
        state_dim,
        policy_config: ContinuousPolicyConfig {
            state_dim,
            hidden_dims: vec![64, 32],
            min_log_std: -5.0,
            max_log_std: 2.0,
            init_log_std: -1.0,
            learnable_std: true,
            action_bounds: (0.0, 1.0),
        },
        value_hidden_dims: vec![64, 32],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        batch_size: 128,
        mini_batch_size: 32,
        num_epochs: 4,
        max_grad_norm: 0.5,
    }
}