foxhunt/crates/ml/examples/hyperopt_baseline_supervised.rs

//! Hyperopt Runner for Supervised Models on DBN Data
//!
//! Runs hyperparameter optimization using Particle Swarm Optimization (PSO) for
//! TFT, MAMBA-2, or both models. Uses the same --data-dir / --symbol interface
//! as the RL hyperopt and training binaries.
//!
//! ## Usage
//!
//! ```bash
//! # Run TFT hyperopt on ES futures
//! hyperopt_baseline_supervised --model tft --data-dir /data/ES.FUT --trials 20 --epochs 20
//!
//! # Run Mamba2 hyperopt
//! hyperopt_baseline_supervised --model mamba2 --data-dir /data/NQ.FUT --trials 20 --epochs 10
//!
//! # Run both models
//! hyperopt_baseline_supervised --model both --data-dir /data/ES.FUT
//! ```
//!
//! ## Output
//!
//! Results are written as JSON to `--output` (default: `ml/trained_models/hyperopt_results.json`).

#![allow(unused_crate_dependencies)]

use anyhow::{Context, Result};
use clap::Parser;
use serde_json::Value;
use std::path::PathBuf;
use std::time::Instant;
use tracing::{error, info, warn, Level};

use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
use ml::hyperopt::adapters::tft::TFTTrainer;
use ml::hyperopt::paths::{generate_run_id, TrainingPaths};
use ml::hyperopt::ArgminOptimizer;

/// Hyperparameter optimization runner for supervised models
#[derive(Parser, Debug)]
#[command(name = "hyperopt-baseline-supervised")]
#[command(about = "Run hyperparameter optimization for supervised models (TFT, Mamba2)")]
struct Args {
    /// Model to optimize: "tft", "mamba2", or "both"
    #[arg(long, default_value = "both")]
    model: String,

    /// Directory containing .dbn.zst OHLCV data files
    #[arg(long)]
    data_dir: PathBuf,

    /// Number of PSO trials per model
    #[arg(long, default_value = "20")]
    trials: usize,

    /// Number of initial LHS (Latin Hypercube Sampling) samples
    #[arg(long, default_value = "5")]
    n_initial: usize,

    /// Training epochs per trial
    #[arg(long, default_value = "20")]
    epochs: usize,

    /// Output path for JSON results
    #[arg(long, default_value = "ml/trained_models/hyperopt_results.json")]
    output: PathBuf,

    /// Random seed for reproducibility
    #[arg(long, default_value = "42")]
    seed: u64,

    /// Base directory for training run outputs (checkpoints, logs, metrics)
    #[arg(long, default_value = "/tmp/ml_training")]
    base_dir: String,

    /// Early stopping patience (epochs without improvement)
    #[arg(long, default_value = "10")]
    early_stopping_patience: usize,
}

fn build_model_result(
    best_objective: f64,
    best_params_json: Value,
    num_trials: usize,
    elapsed_secs: f64,
) -> Value {
    serde_json::json!({
        "best_objective": best_objective,
        "best_params": best_params_json,
        "trials": num_trials,
        "elapsed_secs": elapsed_secs,
    })
}

fn run_tft_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  TFT Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-tft");
    let training_paths = TrainingPaths::new(&args.base_dir, "tft", &run_id);

    info!("Run ID: {}", run_id);
    info!("Data dir: {}", args.data_dir.display());
    info!("Epochs per trial: {}", args.epochs);
    info!("Trials: {}", args.trials);

    let trainer = TFTTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create TFT trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    let start = Instant::now();
    let result = optimizer
        .optimize(trainer)
        .context("TFT hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("TFT hyperopt complete:");
    info!("  Best objective: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());
    info!("  Elapsed: {:.1}s", elapsed);

    let best_params_json = serde_json::to_value(&result.best_params)
        .ok()
        .unwrap_or(Value::Null);

    Ok(build_model_result(
        result.best_objective,
        best_params_json,
        result.all_trials.len(),
        elapsed,
    ))
}

fn run_mamba2_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  Mamba2 Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-mamba2");
    let training_paths = TrainingPaths::new(&args.base_dir, "mamba2", &run_id);

    info!("Run ID: {}", run_id);
    info!("Data dir: {}", args.data_dir.display());
    info!("Epochs per trial: {}", args.epochs);
    info!("Trials: {}", args.trials);

    let trainer = Mamba2Trainer::new(&args.data_dir, args.epochs)
        .context("Failed to create Mamba2 trainer")?
        .with_training_paths(training_paths);

    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    let start = Instant::now();
    let result = optimizer
        .optimize(trainer)
        .context("Mamba2 hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("Mamba2 hyperopt complete:");
    info!("  Best objective: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());
    info!("  Elapsed: {:.1}s", elapsed);

    let best_params_json = serde_json::to_value(&result.best_params)
        .ok()
        .unwrap_or(Value::Null);

    Ok(build_model_result(
        result.best_objective,
        best_params_json,
        result.all_trials.len(),
        elapsed,
    ))
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .with_target(false)
        .init();

    let args = Args::parse();

    info!("========================================");
    info!("  Hyperopt Baseline Supervised Runner");
    info!("========================================");
    info!("Model: {}", args.model);
    info!("Data dir: {}", args.data_dir.display());
    info!("Trials: {}", args.trials);
    info!("Initial LHS samples: {}", args.n_initial);
    info!("Epochs per trial: {}", args.epochs);
    info!("Output: {}", args.output.display());
    info!("Seed: {}", args.seed);

    if !args.data_dir.exists() {
        anyhow::bail!(
            "Data directory not found: {}. Provide a valid path via --data-dir.",
            args.data_dir.display()
        );
    }

    let run_tft = args.model == "tft" || args.model == "both";
    let run_mamba2 = args.model == "mamba2" || args.model == "both";

    if !run_tft && !run_mamba2 {
        anyhow::bail!(
            "Invalid --model value '{}'. Must be 'tft', 'mamba2', or 'both'.",
            args.model
        );
    }

    if args.trials <= args.n_initial {
        anyhow::bail!(
            "trials ({}) must be greater than n_initial ({})",
            args.trials,
            args.n_initial
        );
    }

    if let Some(parent) = args.output.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("Failed to create output directory: {}", parent.display()))?;
    }

    let mut results = serde_json::Map::new();

    if run_tft {
        match run_tft_hyperopt(&args) {
            Ok(tft_result) => {
                results.insert("tft".to_owned(), tft_result);
            }
            Err(e) => {
                error!("TFT hyperopt failed: {:#}", e);
                warn!("Continuing with remaining models...");
                results.insert(
                    "tft".to_owned(),
                    serde_json::json!({ "error": format!("{:#}", e) }),
                );
            }
        }
    }

    if run_mamba2 {
        match run_mamba2_hyperopt(&args) {
            Ok(mamba2_result) => {
                results.insert("mamba2".to_owned(), mamba2_result);
            }
            Err(e) => {
                error!("Mamba2 hyperopt failed: {:#}", e);
                warn!("Continuing...");
                results.insert(
                    "mamba2".to_owned(),
                    serde_json::json!({ "error": format!("{:#}", e) }),
                );
            }
        }
    }

    let output_json = Value::Object(results);
    let output_str =
        serde_json::to_string_pretty(&output_json).context("Failed to serialize results")?;

    std::fs::write(&args.output, &output_str)
        .with_context(|| format!("Failed to write results to {}", args.output.display()))?;

    info!("========================================");
    info!("  Results saved to: {}", args.output.display());
    info!("========================================");
    info!("{}", output_str);

    Ok(())
}