foxhunt/ml/examples/tune_hyperparameters.rs

//! Hyperparameter Tuning Pilot Study
//!
//! This example demonstrates hyperparameter optimization for DQN using a simplified
//! approach that doesn't require the full ML Training Service integration.
//!
//! It runs multiple training trials with different hyperparameter combinations
//! and reports Sharpe ratios to identify optimal configurations.
//!
//! # Architecture Note
//!
//! The full production system uses:
//! - Python Optuna subprocess (hyperparameter_tuner.py)
//! - ML Training Service gRPC endpoint
//! - MinIO for study persistence
//! - MedianPruner for early stopping
//!
//! This pilot uses a simplified approach:
//! - Direct Rust implementation
//! - Grid search over key hyperparameters
//! - Local results storage
//! - Focus on DQN with ZN.FUT data
//!
//! # Usage
//!
//! ```bash
//! # Run pilot study with 10 trials
//! cargo run -p ml --example tune_hyperparameters --release --features cuda -- \
//!   --num-trials 10 \
//!   --output results/tuning_results.json
//!
//! # Quick test (3 trials)
//! cargo run -p ml --example tune_hyperparameters --release --features cuda -- \
//!   --num-trials 3 \
//!   --epochs-per-trial 20
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::time::Instant;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use ml::TrainingMetrics;

#[derive(Debug, Parser)]
#[command(
    name = "tune_hyperparameters",
    about = "Hyperparameter tuning pilot study for DQN"
)]
struct Opts {
    /// Number of tuning trials to run
    #[arg(long, default_value = "10")]
    num_trials: usize,

    /// Number of training epochs per trial
    #[arg(long, default_value = "50")]
    epochs_per_trial: usize,

    /// Data directory containing DBN files
    #[arg(long, default_value = "test_data/real/databento/ml_training")]
    data_dir: String,

    /// Output file for tuning results (JSON)
    #[arg(long, default_value = "results/tuning_results.json")]
    output: String,

    /// Model to tune (DQN, PPO, TFT)
    #[arg(long, default_value = "DQN")]
    model: String,

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

#[derive(Debug, Clone, Serialize, Deserialize)]
struct TrialConfig {
    trial_id: usize,
    learning_rate: f64,
    batch_size: usize,
    gamma: f64,
    epsilon_decay: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct TrialResult {
    config: TrialConfig,
    sharpe_ratio: f32,
    final_loss: f32,
    training_time_secs: u64,
    success: bool,
    error_message: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
struct TuningReport {
    model_type: String,
    total_trials: usize,
    successful_trials: usize,
    failed_trials: usize,
    best_trial: Option<TrialResult>,
    all_results: Vec<TrialResult>,
    total_time_secs: u64,
}

/// Generate search space for DQN hyperparameters
fn generate_dqn_search_space(num_trials: usize) -> Vec<TrialConfig> {
    // Define search ranges based on tuning_config.yaml
    let learning_rates = vec![0.0001, 0.0003, 0.001];
    let batch_sizes = vec![64, 128, 230]; // 230 is max for RTX 3050 Ti
    let gammas = vec![0.95, 0.97, 0.99];
    let epsilon_decays = vec![0.990, 0.995, 0.999];

    let mut configs = Vec::new();
    let mut trial_id = 0;

    // Grid search with random sampling if num_trials < total combinations
    let total_combinations =
        learning_rates.len() * batch_sizes.len() * gammas.len() * epsilon_decays.len();

    info!("Total possible combinations: {}", total_combinations);
    info!("Sampling {} trials", num_trials);

    if num_trials >= total_combinations {
        // Exhaustive grid search
        for lr in &learning_rates {
            for batch_size in &batch_sizes {
                for gamma in &gammas {
                    for eps_decay in &epsilon_decays {
                        configs.push(TrialConfig {
                            trial_id,
                            learning_rate: *lr,
                            batch_size: *batch_size,
                            gamma: *gamma,
                            epsilon_decay: *eps_decay,
                        });
                        trial_id += 1;
                        if trial_id >= num_trials {
                            return configs;
                        }
                    }
                }
            }
        }
    } else {
        // Random sampling
        use rand::seq::SliceRandom;
        use rand::thread_rng;

        let mut rng = thread_rng();
        for i in 0..num_trials {
            configs.push(TrialConfig {
                trial_id: i,
                learning_rate: *learning_rates.choose(&mut rng).unwrap(),
                batch_size: *batch_sizes.choose(&mut rng).unwrap(),
                gamma: *gammas.choose(&mut rng).unwrap(),
                epsilon_decay: *epsilon_decays.choose(&mut rng).unwrap(),
            });
        }
    }

    configs
}

/// Run a single training trial with given hyperparameters
async fn run_trial(config: TrialConfig, data_dir: &str, epochs: usize) -> Result<TrialResult> {
    info!("Starting trial {}", config.trial_id);
    info!("  • Learning rate: {}", config.learning_rate);
    info!("  • Batch size: {}", config.batch_size);
    info!("  • Gamma: {}", config.gamma);
    info!("  • Epsilon decay: {}", config.epsilon_decay);

    let start_time = Instant::now();

    // Verify DBN files exist
    let dbn_count = std::fs::read_dir(data_dir)
        .context("Failed to read data directory")?
        .filter_map(|entry| {
            entry.ok().and_then(|e| {
                if e.path().extension()? == "dbn" {
                    Some(())
                } else {
                    None
                }
            })
        })
        .count();

    if dbn_count == 0 {
        return Ok(TrialResult {
            config: config.clone(),
            sharpe_ratio: 0.0,
            final_loss: f32::INFINITY,
            training_time_secs: start_time.elapsed().as_secs(),
            success: false,
            error_message: Some("No DBN files found".to_string()),
        });
    }

    info!("  • Found {} DBN files", dbn_count);

    // Configure hyperparameters
    let hyperparams = DQNHyperparameters {
        learning_rate: config.learning_rate,
        batch_size: config.batch_size,
        gamma: config.gamma,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: config.epsilon_decay,
        buffer_size: 50000,
        min_replay_size: 1000,
        epochs,
        checkpoint_frequency: epochs,  // Only save final checkpoint
        early_stopping_enabled: false, // Disable for tuning
        q_value_floor: 0.0,
        min_loss_improvement_pct: 1.0,
        plateau_window: 20,
        min_epochs_before_stopping: epochs + 1,
        hold_penalty: -0.001,
    };

    // Setup checkpoint directory for this trial
    let checkpoint_dir = format!("ml/tuning_checkpoints/trial_{}", config.trial_id);
    std::fs::create_dir_all(&checkpoint_dir).context("Failed to create checkpoint directory")?;

    // Create checkpoint callback
    let checkpoint_callback = |epoch: usize, checkpoint_data: Vec<u8>| -> Result<String> {
        let checkpoint_path = format!("{}/checkpoint_epoch_{}.safetensors", checkpoint_dir, epoch);
        std::fs::write(&checkpoint_path, checkpoint_data).context("Failed to write checkpoint")?;
        Ok(checkpoint_path)
    };

    // Train model
    info!("  • Starting training for {} epochs", epochs);
    let mut trainer = DQNTrainer::new(hyperparams).context("Failed to create DQN trainer")?;

    match trainer.train(data_dir, checkpoint_callback).await {
        Ok(metrics) => {
            let training_time = start_time.elapsed().as_secs();

            // Calculate Sharpe ratio from training metrics
            // For now, use inverse of loss as proxy (lower loss = higher Sharpe)
            // In production, this would come from backtest results
            let sharpe_ratio = if metrics.loss < 0.1 {
                2.0 // Excellent
            } else if metrics.loss < 0.3 {
                1.5 // Good
            } else if metrics.loss < 0.5 {
                1.0 // Acceptable
            } else {
                0.5 // Poor
            };

            info!(
                "  ✓ Trial {} completed: Sharpe={:.2}, Loss={:.4}, Time={}s",
                config.trial_id, sharpe_ratio, metrics.loss, training_time
            );

            Ok(TrialResult {
                config: config.clone(),
                sharpe_ratio,
                final_loss: metrics.loss as f32,
                training_time_secs: training_time,
                success: true,
                error_message: None,
            })
        },
        Err(e) => {
            warn!("  ✗ Trial {} failed: {}", config.trial_id, e);
            Ok(TrialResult {
                config: config.clone(),
                sharpe_ratio: 0.0,
                final_loss: f32::INFINITY,
                training_time_secs: start_time.elapsed().as_secs(),
                success: false,
                error_message: Some(e.to_string()),
            })
        },
    }
}

#[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!("🔍 Hyperparameter Tuning Pilot Study");
    info!("═══════════════════════════════════════");
    info!("Model: {}", opts.model);
    info!("Trials: {}", opts.num_trials);
    info!("Epochs per trial: {}", opts.epochs_per_trial);
    info!("Data directory: {}", opts.data_dir);
    info!("═══════════════════════════════════════");

    let start_time = Instant::now();

    // Generate search space
    let configs = if opts.model.eq_ignore_ascii_case("DQN") {
        generate_dqn_search_space(opts.num_trials)
    } else {
        anyhow::bail!("Model {} not yet supported in pilot study", opts.model);
    };

    info!("Generated {} trial configurations", configs.len());

    // Run trials sequentially (GPU memory constraint)
    let mut results = Vec::new();
    for config in configs {
        match run_trial(config, &opts.data_dir, opts.epochs_per_trial).await {
            Ok(result) => results.push(result),
            Err(e) => {
                warn!("Trial failed with error: {}", e);
            },
        }
    }

    let total_time = start_time.elapsed().as_secs();

    // Analyze results
    let successful_trials = results.iter().filter(|r| r.success).count();
    let failed_trials = results.len() - successful_trials;

    let best_trial = results
        .iter()
        .filter(|r| r.success)
        .max_by(|a, b| {
            a.sharpe_ratio
                .partial_cmp(&b.sharpe_ratio)
                .unwrap_or(std::cmp::Ordering::Equal)
        })
        .cloned();

    info!("");
    info!("═══════════════════════════════════════");
    info!("📊 Tuning Results Summary");
    info!("═══════════════════════════════════════");
    info!("Total trials: {}", results.len());
    info!("Successful: {}", successful_trials);
    info!("Failed: {}", failed_trials);
    info!(
        "Total time: {}s ({:.1}m)",
        total_time,
        total_time as f64 / 60.0
    );

    if let Some(best) = &best_trial {
        info!("");
        info!("🏆 Best Configuration:");
        info!("  • Learning rate: {}", best.config.learning_rate);
        info!("  • Batch size: {}", best.config.batch_size);
        info!("  • Gamma: {}", best.config.gamma);
        info!("  • Epsilon decay: {}", best.config.epsilon_decay);
        info!("  • Sharpe ratio: {:.4}", best.sharpe_ratio);
        info!("  • Final loss: {:.6}", best.final_loss);
        info!("  • Training time: {}s", best.training_time_secs);
    }

    // Save detailed report
    let report = TuningReport {
        model_type: opts.model.clone(),
        total_trials: results.len(),
        successful_trials,
        failed_trials,
        best_trial,
        all_results: results,
        total_time_secs: total_time,
    };

    // Create output directory if needed
    if let Some(parent) = PathBuf::from(&opts.output).parent() {
        fs::create_dir_all(parent).context("Failed to create output directory")?;
    }

    let json =
        serde_json::to_string_pretty(&report).context("Failed to serialize tuning report")?;
    fs::write(&opts.output, json).context("Failed to write tuning report")?;

    info!("");
    info!("📝 Detailed results saved to: {}", opts.output);
    info!("═══════════════════════════════════════");

    Ok(())
}