foxhunt/ml/examples/optimize_mamba2_standalone.rs

//! Standalone hyperparameter optimization for MAMBA-2
//!
//! This is a complete, self-contained example for Bayesian hyperparameter optimization
//! of the MAMBA-2 State Space Model. It uses the egobox library for efficient
//! Gaussian Process-based optimization with Expected Improvement acquisition.
//!
//! # Features
//!
//! - **Bayesian Optimization**: Efficiently finds optimal hyperparameters in 20-30 trials
//! - **GPU Accelerated**: Each trial runs on CUDA for fast evaluation
//! - **Latin Hypercube Sampling**: Smart initialization for exploration
//! - **Progress Tracking**: Real-time updates with trial metrics
//! - **YAML Export**: Save results for production deployment
//! - **ASCII Convergence Plot**: Visual feedback on optimization progress
//!
//! # Search Space
//!
//! The optimizer searches over 4 hyperparameters:
//! - Learning rate: 1e-5 to 1e-2 (log scale)
//! - Batch size: 16 to 256 (integer, discrete)
//! - Dropout: 0.0 to 0.5 (linear scale)
//! - Weight decay: 1e-6 to 1e-2 (log scale)
//!
//! # Performance
//!
//! - Trial duration: ~18 seconds (10 epochs)
//! - Total time (30 trials): ~9 minutes
//! - GPU memory: ~2GB VRAM per trial
//! - Cost (RTX A4000): ~$0.04 (9 min @ $0.25/hr)
//!
//! # Usage
//!
//! ```bash
//! # Default: 30 trials on ES.FUT data
//! cargo run -p ml --example optimize_mamba2_standalone --release --features cuda
//!
//! # Custom configuration:
//! cargo run -p ml --example optimize_mamba2_standalone --release --features cuda -- \
//!     --parquet-file test_data/NQ_FUT_180d.parquet \
//!     --max-trials 50 \
//!     --epochs-per-trial 15 \
//!     --output best_mamba2_params.yaml
//!
//! # Show help:
//! cargo run -p ml --example optimize_mamba2_standalone --release --features cuda -- --help
//! ```
//!
//! # Output
//!
//! The script produces:
//! 1. Console output with trial-by-trial progress
//! 2. YAML file with best hyperparameters (if --output specified)
//! 3. ASCII convergence plot showing optimization progress
//! 4. Summary statistics and next steps
//!
//! # Example Output
//!
//! ```text
//! ╔═══════════════════════════════════════════════════════════╗
//! ║  MAMBA-2 Bayesian Hyperparameter Optimization            ║
//! ╚═══════════════════════════════════════════════════════════╝
//!
//! Configuration:
//!   Parquet File: test_data/ES_FUT_180d.parquet
//!   Max Trials: 30
//!   Epochs per Trial: 10
//!   Search Space:
//!     Learning Rate: 10^-5.0 to 10^-2.0
//!     Batch Size: 16 to 256
//!     Dropout: 0.00 to 0.50
//!     Weight Decay: 10^-6.0 to 10^-2.0
//!
//! ╔═══════════════════════════════════════════════════════════╗
//! ║  Trial 1: Evaluating Hyperparameters                     ║
//! ╚═══════════════════════════════════════════════════════════╝
//!   Learning Rate: 0.000543
//!   Batch Size: 128
//!   Dropout: 0.234
//!   Weight Decay: 0.000089
//! ✓ Trial 1 completed in 18.2s
//!   Validation Loss: 0.123456
//!   Perplexity: 1.1315
//!
//! [... 28 more trials ...]
//!
//! ╔═══════════════════════════════════════════════════════════╗
//! ║  Optimization Complete                                    ║
//! ╚═══════════════════════════════════════════════════════════╝
//!
//! Best Hyperparameters Found:
//!   Learning Rate: 0.000321
//!   Batch Size: 64
//!   Dropout: 0.150
//!   Weight Decay: 0.000045
//!   Best Validation Loss: 0.098765
//!   Best Perplexity: 1.1037
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use tracing::info;

use ml::hyperopt::egobox_tuner::{optimize_mamba2, HyperparameterSpace, OptimizationResult};

/// CLI arguments for MAMBA-2 hyperparameter optimization
#[derive(Parser, Debug)]
#[command(
    name = "optimize_mamba2_standalone",
    about = "Standalone MAMBA-2 Bayesian Hyperparameter Optimization",
    long_about = "Complete, self-contained example for optimizing MAMBA-2 hyperparameters using Bayesian optimization with Gaussian Process surrogates and Expected Improvement acquisition."
)]
struct Args {
    /// Path to Parquet file with training data
    #[arg(
        long,
        default_value = "test_data/ES_FUT_180d.parquet",
        help = "Parquet file containing OHLCV market data"
    )]
    parquet_file: PathBuf,

    /// Maximum number of trials
    #[arg(
        long,
        default_value = "30",
        help = "Total optimization trials (includes 5 initial LHS samples)"
    )]
    max_trials: usize,

    /// Output YAML file for best parameters
    #[arg(
        long,
        default_value = "best_params.yaml",
        help = "Output file for best hyperparameters (YAML format)"
    )]
    output: PathBuf,

    /// Number of epochs per trial (shorter = faster)
    #[arg(
        long,
        default_value = "10",
        help = "Training epochs per trial (lower for faster feedback)"
    )]
    epochs_per_trial: usize,

    /// Minimum learning rate (actual value, not log)
    #[arg(long, default_value = "0.00001", help = "Minimum learning rate (1e-5)")]
    lr_min: f64,

    /// Maximum learning rate (actual value, not log)
    #[arg(long, default_value = "0.01", help = "Maximum learning rate (1e-2)")]
    lr_max: f64,

    /// Minimum batch size
    #[arg(long, default_value = "16", help = "Minimum batch size")]
    batch_size_min: usize,

    /// Maximum batch size
    #[arg(long, default_value = "256", help = "Maximum batch size")]
    batch_size_max: usize,

    /// Minimum dropout rate
    #[arg(
        long,
        default_value = "0.0",
        help = "Minimum dropout (0.0 = no dropout)"
    )]
    dropout_min: f64,

    /// Maximum dropout rate
    #[arg(
        long,
        default_value = "0.5",
        help = "Maximum dropout (0.5 = aggressive)"
    )]
    dropout_max: f64,

    /// Minimum weight decay (actual value, not log)
    #[arg(long, default_value = "0.000001", help = "Minimum weight decay (1e-6)")]
    wd_min: f64,

    /// Maximum weight decay (actual value, not log)
    #[arg(long, default_value = "0.01", help = "Maximum weight decay (1e-2)")]
    wd_max: f64,

    /// Skip convergence plot
    #[arg(
        long,
        default_value = "false",
        help = "Skip ASCII convergence plot generation"
    )]
    no_plot: bool,
}

impl Args {
    /// Validate CLI arguments
    fn validate(&self) -> Result<()> {
        // Validate learning rate bounds
        if self.lr_min <= 0.0 || self.lr_min >= self.lr_max {
            anyhow::bail!(
                "Invalid learning rate: min={}, max={}. Must be 0 < min < max",
                self.lr_min,
                self.lr_max
            );
        }

        // Validate weight decay bounds
        if self.wd_min < 0.0 || self.wd_min >= self.wd_max {
            anyhow::bail!(
                "Invalid weight decay: min={}, max={}. Must be 0 <= min < max",
                self.wd_min,
                self.wd_max
            );
        }

        // Validate dropout bounds
        if self.dropout_min < 0.0 || self.dropout_min >= self.dropout_max || self.dropout_max > 1.0
        {
            anyhow::bail!(
                "Invalid dropout: min={}, max={}. Must be 0 <= min < max <= 1",
                self.dropout_min,
                self.dropout_max
            );
        }

        // Validate batch size bounds
        if self.batch_size_min == 0 || self.batch_size_min >= self.batch_size_max {
            anyhow::bail!(
                "Invalid batch size: min={}, max={}. Must be 0 < min < max",
                self.batch_size_min,
                self.batch_size_max
            );
        }

        // Validate trials
        if self.max_trials < 6 {
            anyhow::bail!("Max trials must be >= 6 (5 initial LHS samples + 1 optimization)");
        }

        // Validate epochs
        if self.epochs_per_trial == 0 {
            anyhow::bail!("Epochs per trial must be > 0");
        }

        // Validate Parquet file exists
        if !self.parquet_file.exists() {
            anyhow::bail!("Parquet file not found: {:?}", self.parquet_file);
        }

        Ok(())
    }

    /// Convert to HyperparameterSpace
    fn to_space(&self) -> HyperparameterSpace {
        HyperparameterSpace {
            learning_rate_log_min: self.lr_min.log10(),
            learning_rate_log_max: self.lr_max.log10(),
            batch_size_min: self.batch_size_min,
            batch_size_max: self.batch_size_max,
            dropout_min: self.dropout_min,
            dropout_max: self.dropout_max,
            weight_decay_log_min: self.wd_min.log10(),
            weight_decay_log_max: self.wd_max.log10(),
        }
    }
}

/// Generate ASCII convergence plot
fn generate_convergence_plot(result: &OptimizationResult) -> String {
    let trials = &result.trial_history;
    if trials.is_empty() {
        return "No trials to plot".to_string();
    }

    // Find min/max losses for scaling
    let min_loss = trials
        .iter()
        .map(|t| t.validation_loss)
        .fold(f64::INFINITY, f64::min);
    let max_loss = trials
        .iter()
        .map(|t| t.validation_loss)
        .fold(f64::NEG_INFINITY, f64::max);

    let mut plot = String::new();
    plot.push_str("Convergence Plot (Validation Loss vs Trial)\n\n");

    // ASCII plot dimensions
    let height = 20;
    let width = 60;

    // Scale losses to plot height
    let scale = |loss: f64| -> usize {
        let normalized = (loss - min_loss) / (max_loss - min_loss).max(1e-6);
        height - ((normalized * (height as f64)).round() as usize).min(height - 1)
    };

    // Create plot grid
    let mut grid = vec![vec![' '; width]; height];

    // Plot points
    for (i, trial) in trials.iter().enumerate() {
        let x = ((i as f64 / trials.len().max(1) as f64) * (width - 1) as f64).round() as usize;
        let y = scale(trial.validation_loss);
        grid[y][x] = '*';
    }

    // Add Y-axis
    plot.push_str(&format!("{:.4} |", max_loss));
    for _ in 0..width {
        plot.push('-');
    }
    plot.push('\n');

    for row in &grid {
        plot.push_str("       |");
        for &ch in row {
            plot.push(ch);
        }
        plot.push('\n');
    }

    plot.push_str(&format!("{:.4} |", min_loss));
    for _ in 0..width {
        plot.push('-');
    }
    plot.push('\n');

    plot.push_str("       0");
    for _ in 0..(width - 10) {
        plot.push(' ');
    }
    plot.push_str(&format!("{}\n", trials.len()));

    plot
}

/// Main optimization function
#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_target(false)
        .with_thread_ids(false)
        .init();

    // Parse and validate arguments
    let args = Args::parse();

    if let Err(e) = args.validate() {
        tracing::error!("Invalid arguments: {}", e);
        std::process::exit(1);
    }

    // Create hyperparameter space
    let space = args.to_space();

    // Run optimization
    info!("");
    let estimated_duration = (args.max_trials as f64 * args.epochs_per_trial as f64 * 1.8) / 60.0;
    info!("Expected duration: ~{:.1} minutes", estimated_duration);
    info!("");

    let result = optimize_mamba2(
        space,
        args.parquet_file.to_str().unwrap(),
        args.max_trials,
        args.epochs_per_trial,
    )
    .await
    .context("Optimization failed")?;

    // Display results
    print_results(&result);

    // Generate convergence plot
    if !args.no_plot && !result.trial_history.is_empty() {
        info!("");
        let plot = generate_convergence_plot(&result);
        println!("{}", plot);
    }

    // Save to YAML
    save_results(&args.output, &result)?;

    // Print next steps
    print_next_steps(&result);

    Ok(())
}

/// Print optimization results
fn print_results(result: &OptimizationResult) {
    info!("");
    info!("╔═══════════════════════════════════════════════════════════╗");
    info!("║   Optimization Results                                    ║");
    info!("╚═══════════════════════════════════════════════════════════╝");
    info!("Best Hyperparameters:");
    info!("  Learning Rate: {:.6}", result.best_params.learning_rate);
    info!("  Batch Size: {}", result.best_params.batch_size);
    info!("  Dropout: {:.3}", result.best_params.dropout);
    info!("  Weight Decay: {:.6}", result.best_params.weight_decay);
    info!("");
    info!("Performance:");
    info!(
        "  Best Validation Loss: {:.6}",
        result.best_params.best_validation_loss
    );
    info!(
        "  Best Perplexity: {:.4}",
        result.best_params.best_validation_loss.exp()
    );
    info!("  Trials Used: {}", result.best_params.trials_used);
}

/// Save results to YAML
fn save_results(output_file: &PathBuf, result: &OptimizationResult) -> Result<()> {
    // Create parent directories if needed
    if let Some(parent) = output_file.parent() {
        std::fs::create_dir_all(parent).context("Failed to create output directory")?;
    }

    let yaml_content =
        serde_yaml::to_string(&result.best_params).context("Failed to serialize to YAML")?;

    std::fs::write(output_file, yaml_content).context("Failed to write YAML file")?;

    info!("");
    info!("✓ Best hyperparameters saved to: {:?}", output_file);

    Ok(())
}

/// Print next steps
fn print_next_steps(result: &OptimizationResult) {
    info!("");
    info!("╔═══════════════════════════════════════════════════════════╗");
    info!("║   Next Steps                                              ║");
    info!("╚═══════════════════════════════════════════════════════════╝");
    info!("1. Full Training (50-200 epochs):");
    info!("   cargo run -p ml --example train_mamba2_parquet --release --features cuda -- \\");
    info!(
        "       --learning-rate {} \\",
        result.best_params.learning_rate
    );
    info!("       --batch-size {} \\", result.best_params.batch_size);
    info!("       --dropout {} \\", result.best_params.dropout);
    info!(
        "       --weight-decay {} \\",
        result.best_params.weight_decay
    );
    info!("       --epochs 100");
    info!("");
    info!("2. Deploy to Runpod GPU:");
    info!("   python3 scripts/runpod_deploy.py --gpu-type \"RTX A4000\" \\");
    info!("       --training-script train_mamba2_parquet \\");
    info!(
        "       --extra-args \"--learning-rate {} --batch-size {} --dropout {} --weight-decay {}\"",
        result.best_params.learning_rate,
        result.best_params.batch_size,
        result.best_params.dropout,
        result.best_params.weight_decay
    );
    info!("");
    info!("3. Production Deployment:");
    info!("   - Save model to S3: s3://se3zdnb5o4/models/mamba2_optimized.safetensors");
    info!("   - Update model registry: ml/models/registry.yaml");
    info!("   - Run A/B test: cargo test --package ml --test model_ab_test");
}