foxhunt/ml/examples/train_ppo_parquet.rs

//! PPO Training Example with Parquet Data
//!
//! Trains a PPO model on market data from Parquet files with:
//! - Real OHLCV data + 51-dimensional features
//! - Actual PnL-based rewards
//! - GAE advantages on real price trajectories
//! - Policy convergence validation (KL divergence > 0)
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (30 epochs, hyperopt-optimized learning rates)
//! cargo run -p ml --example train_ppo_parquet --release --features cuda -- \
//!   --parquet-file test_data/ZN_FUT_90d_clean.parquet
//!
//! # Custom epochs, batch size, and learning rates
//! cargo run -p ml --example train_ppo_parquet --release --features cuda -- \
//!   --parquet-file test_data/ZN_FUT_90d_clean.parquet \
//!   --epochs 50 \
//!   --batch-size 128 \
//!   --policy-lr 0.000001 \
//!   --value-lr 0.001
//!
//! # With early stopping disabled
//! cargo run -p ml --example train_ppo_parquet --release --features cuda -- \
//!   --parquet-file test_data/NQ_FUT_180d.parquet \
//!   --no-early-stopping
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use std::fs::File;
use std::path::PathBuf;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

use arrow::array::{Array, Float64Array, PrimitiveArray, UInt64Array};
use arrow::datatypes::TimestampNanosecondType;
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

use ml::features::extraction::{extract_ml_features, OHLCVBar};
use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer, PpoTrainingMetrics};

/// Train PPO model on Parquet market data
#[derive(Debug, Parser)]
#[command(
    name = "train_ppo_parquet",
    about = "Train PPO model on Parquet market data"
)]
struct Opts {
    /// Path to Parquet file with market data
    #[arg(long)]
    parquet_file: String,

    /// Number of training epochs (default: 30 for policy convergence)
    #[arg(long, default_value = "30")]
    epochs: usize,

    /// Policy (actor) learning rate (default: 1e-6, ultra-conservative for stability)
    #[arg(long, default_value = "0.000001")]
    policy_lr: f64,

    /// Value (critic) learning rate (default: 0.001, aggressive for faster convergence)
    #[arg(long, default_value = "0.001")]
    value_lr: f64,

    /// Batch size (max 230 for RTX 3050 Ti 4GB)
    #[arg(long, default_value = "64")]
    batch_size: usize,

    /// Output directory for trained model
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// Verbose logging
    #[arg(short, long)]
    verbose: bool,

    /// Enable early stopping (recommended, use --no-early-stopping to disable)
    #[arg(long)]
    early_stopping: bool,

    /// Disable early stopping
    #[arg(long)]
    no_early_stopping: bool,

    /// Minimum value loss improvement percentage for plateau detection
    #[arg(long, default_value = "2.0")]
    min_value_loss_improvement: f64,

    /// Minimum explained variance threshold
    #[arg(long, default_value = "0.4")]
    min_explained_variance: f64,

    /// Plateau detection window size (epochs)
    #[arg(long, default_value = "30")]
    plateau_window: usize,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Parse CLI options
    let opts = Opts::parse();

    // Setup logging
    let level = if opts.verbose {
        tracing::Level::DEBUG
    } else {
        tracing::Level::INFO
    };

    let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    info!("🚀 Starting PPO Training with Parquet Data");
    info!("Configuration:");
    info!("  • Parquet file: {}", opts.parquet_file);
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Policy learning rate: {}", opts.policy_lr);
    info!("  • Value learning rate: {}", opts.value_lr);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • GPU: CUDA if available (auto-fallback to CPU)");
    info!("  • Output directory: {}", opts.output_dir);

    // Determine early stopping (enabled by default, unless --no-early-stopping is specified)
    let early_stopping_enabled = !opts.no_early_stopping;
    info!(
        "  • Early stopping: {}",
        if early_stopping_enabled {
            "enabled"
        } else {
            "disabled"
        }
    );
    if early_stopping_enabled {
        info!(
            "    - Min value loss improvement: {}%",
            opts.min_value_loss_improvement
        );
        info!(
            "    - Min explained variance: {}",
            opts.min_explained_variance
        );
        info!("    - Plateau window: {} epochs", opts.plateau_window);
    }

    // Create output directory
    let output_path = PathBuf::from(&opts.output_dir);
    if !output_path.exists() {
        std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
        info!("✅ Created output directory: {}", opts.output_dir);
    }

    // Load market data from Parquet file
    info!("\n📊 Loading market data from Parquet file...");
    let bars = load_parquet_data(&opts.parquet_file)
        .await
        .context("Failed to load Parquet data")?;

    info!("✅ Loaded {} OHLCV bars", bars.len());

    // Extract 51-dimensional feature vectors (43 base + 8 OFI)
    info!("\n🏗️  Extracting 51-dimensional feature vectors...");
    let feature_vectors =
        extract_ml_features(&bars).context("Failed to extract 51-dimensional features")?;

    info!(
        "✅ Extracted {} feature vectors (dim=51, warmup bars skipped=50)",
        feature_vectors.len()
    );

    // Convert FeatureVector ([f64; 51]) to Vec<Vec<f32>> for PPO trainer
    let state_dim = 51;
    let market_data: Vec<Vec<f32>> = feature_vectors
        .iter()
        .map(|fv| fv.iter().map(|&v| v as f32).collect())
        .collect();

    // Validate state dimensions
    if let Some(first_state) = market_data.first() {
        if first_state.len() != state_dim {
            return Err(anyhow::anyhow!(
                "State dimension mismatch: expected {}, got {}",
                state_dim,
                first_state.len()
            ));
        }
    }

    info!(
        "✅ Feature extraction complete: {} samples",
        market_data.len()
    );

    // Configure PPO hyperparameters
    let hyperparams = PpoHyperparameters {
        learning_rate: 1e-4, // Deprecated field, kept for backward compatibility
        actor_learning_rate: Some(opts.policy_lr),
        critic_learning_rate: Some(opts.value_lr),
        batch_size: opts.batch_size,
        gamma: 0.99,
        clip_epsilon: 0.2,
        vf_coef: 0.5,
        ent_coef: 0.01,
        gae_lambda: 0.95,
        rollout_steps: 2048,
        minibatch_size: opts.batch_size,
        epochs: opts.epochs,
        early_stopping_enabled,
        min_value_loss_improvement_pct: opts.min_value_loss_improvement,
        min_explained_variance: opts.min_explained_variance,
        plateau_window: opts.plateau_window,
        min_epochs_before_stopping: 50,
        ..PpoHyperparameters::conservative()
    };

    // Create PPO trainer
    let trainer = PpoTrainer::new(
        hyperparams.clone(),
        state_dim,
        &opts.output_dir,
        true, // Use GPU if available
        None, // Single environment (standard mode)
    )
    .context("Failed to create PPO trainer")?;

    info!("✅ PPO trainer initialized (state_dim={})", state_dim);

    // Create progress callback with convergence tracking
    let mut policy_updates = 0;
    let mut kl_divergence_history = Vec::new();

    let progress_callback = |metrics: PpoTrainingMetrics| {
        // Track policy updates (KL divergence > 0 indicates policy changed)
        if metrics.kl_divergence > 0.0 {
            policy_updates += 1;
        }
        kl_divergence_history.push(metrics.kl_divergence);

        info!(
            "📊 Epoch {}/{}: policy_loss={:.4}, value_loss={:.4}, kl_div={:.6}, expl_var={:.4}, mean_reward={:.4}",
            metrics.epoch,
            hyperparams.epochs,
            metrics.policy_loss,
            metrics.value_loss,
            metrics.kl_divergence,
            metrics.explained_variance,
            metrics.mean_reward
        );
    };

    // Train the model
    info!("\n🏋️  Starting training...\n");
    let start_time = std::time::Instant::now();

    let final_metrics = trainer
        .train(market_data, progress_callback)
        .await
        .context("Training failed")?;

    let training_duration = start_time.elapsed();

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    info!("  • Policy loss: {:.6}", final_metrics.policy_loss);
    info!("  • Value loss: {:.6}", final_metrics.value_loss);
    info!("  • KL divergence: {:.6}", final_metrics.kl_divergence);
    info!(
        "  • Explained variance: {:.4}",
        final_metrics.explained_variance
    );
    info!("  • Mean reward: {:.4}", final_metrics.mean_reward);
    info!("  • Std reward: {:.4}", final_metrics.std_reward);
    info!("  • Entropy: {:.4}", final_metrics.entropy);
    info!(
        "  • Training time: {:.1}s ({:.1} min)",
        training_duration.as_secs_f64(),
        training_duration.as_secs_f64() / 60.0
    );

    // Validate policy convergence
    info!("\n🔍 Policy Convergence Analysis:");
    info!("  • Total epochs: {}", hyperparams.epochs);
    info!("  • Policy updates (KL > 0): {}", policy_updates);
    info!(
        "  • Policy update rate: {:.1}%",
        (policy_updates as f64 / hyperparams.epochs as f64) * 100.0
    );

    // Calculate KL divergence statistics
    let kl_mean = kl_divergence_history.iter().sum::<f32>() / kl_divergence_history.len() as f32;
    let kl_max = kl_divergence_history
        .iter()
        .copied()
        .fold(f32::NEG_INFINITY, f32::max);
    let kl_min = kl_divergence_history
        .iter()
        .copied()
        .fold(f32::INFINITY, f32::min);

    info!("  • KL divergence (mean): {:.6}", kl_mean);
    info!("  • KL divergence (max): {:.6}", kl_max);
    info!("  • KL divergence (min): {:.6}", kl_min);

    // Convergence validation
    if final_metrics.kl_divergence > 0.0 {
        info!("  ✅ PASS: Policy updates detected (KL divergence > 0)");
    } else {
        warn!("  ⚠️  WARN: No policy updates in final epoch (KL divergence = 0)");
        warn!("          This may indicate learning rate too low or convergence");
    }

    // Value function validation
    if final_metrics.explained_variance > 0.5 {
        info!("  ✅ PASS: Value network learning (explained variance > 0.5)");
    } else {
        warn!("  ⚠️  WARN: Value network may need tuning (explained variance < 0.5)");
    }

    // Checkpoint is already saved by trainer (every 10 epochs)
    let final_checkpoint = output_path.join(format!(
        "ppo_checkpoint_epoch_{}.safetensors",
        hyperparams.epochs
    ));

    info!(
        "\n💾 Final checkpoint saved to: {}",
        final_checkpoint.display()
    );
    info!("\n🎉 PPO training complete with Parquet data!");
    info!("📁 Model files saved to: {}", opts.output_dir);

    info!("\n📈 Training Summary:");
    info!("  • Data source: Parquet file ({})", opts.parquet_file);
    info!("  • Training samples: {}", bars.len());
    info!(
        "  • Feature samples: {} (after warmup)",
        feature_vectors.len()
    );
    info!("  • State dimension: {}", state_dim);
    info!("  • Features: 51-dimensional");
    info!(
        "  • Policy updates: {}/{} epochs ({:.1}%)",
        policy_updates,
        hyperparams.epochs,
        (policy_updates as f64 / hyperparams.epochs as f64) * 100.0
    );
    info!(
        "  • Convergence: {}",
        if final_metrics.kl_divergence > 0.0 {
            "✅ Achieved"
        } else {
            "⚠️  Check logs"
        }
    );

    Ok(())
}

/// Load OHLCV data from Parquet file (Databento schema)
async fn load_parquet_data(parquet_path: &str) -> Result<Vec<OHLCVBar>> {
    info!("Loading Parquet file: {}", parquet_path);

    // Open Parquet file
    let file = File::open(parquet_path)
        .with_context(|| format!("Failed to open Parquet file: {}", parquet_path))?;

    // Create Parquet reader
    let builder = ParquetRecordBatchReaderBuilder::try_new(file)
        .with_context(|| "Failed to create Parquet reader")?;

    let reader = builder
        .build()
        .with_context(|| "Failed to build Parquet reader")?;

    // Read all batches
    let mut all_ohlcv_bars = Vec::new();

    for batch_result in reader {
        let batch: RecordBatch = batch_result.with_context(|| "Failed to read record batch")?;

        // Extract columns from Databento Parquet schema:
        // Column 3: open, Column 4: high, Column 5: low, Column 6: close
        // Column 7: volume, Column 9: ts_event (Timestamp(Nanosecond, Some("UTC")))
        let timestamps = batch
            .column(9)
            .as_any()
            .downcast_ref::<PrimitiveArray<TimestampNanosecondType>>()
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "Failed to downcast timestamp column. Expected Timestamp(Nanosecond), got: {:?}",
                    batch.column(9).data_type()
                )
            })?;

        let opens = batch
            .column(3)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| anyhow::anyhow!("Failed to downcast open column"))?;

        let highs = batch
            .column(4)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| anyhow::anyhow!("Failed to downcast high column"))?;

        let lows = batch
            .column(5)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| anyhow::anyhow!("Failed to downcast low column"))?;

        let closes = batch
            .column(6)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| anyhow::anyhow!("Failed to downcast close column"))?;

        let volumes = batch
            .column(7)
            .as_any()
            .downcast_ref::<UInt64Array>()
            .ok_or_else(|| anyhow::anyhow!("Failed to downcast volume column"))?;

        // Convert to OHLCVBar structs
        for i in 0..batch.num_rows() {
            let timestamp_ns = timestamps.value(i);

            // Convert nanoseconds to DateTime<Utc>
            let timestamp = chrono::DateTime::from_timestamp(
                (timestamp_ns / 1_000_000_000) as i64,
                (timestamp_ns % 1_000_000_000) as u32,
            )
            .unwrap_or_else(|| chrono::Utc::now());

            let bar = OHLCVBar {
                timestamp,
                open: opens.value(i),
                high: highs.value(i),
                low: lows.value(i),
                close: closes.value(i),
                volume: volumes.value(i) as f64,
            };

            all_ohlcv_bars.push(bar);
        }
    }

    info!("✅ Loaded {} OHLCV bars from Parquet", all_ohlcv_bars.len());

    Ok(all_ohlcv_bars)
}