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
//! any of the 8 supervised models: TFT, Mamba2, Liquid, TGGN, TLOB, KAN, xLSTM,
//! and Diffusion. 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 a specific model
//! hyperopt_baseline_supervised --model liquid --data-dir /data/ES.FUT
//!
//! # Run TFT + Mamba2 (backward-compatible alias)
//! hyperopt_baseline_supervised --model both --data-dir /data/ES.FUT
//!
//! # Run all 8 supervised models
//! hyperopt_baseline_supervised --model all --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};

use common::metrics::{server as metrics_server, training_metrics};

use ml::hyperopt::adapters::diffusion::DiffusionTrainer;
use ml::hyperopt::adapters::kan::KANTrainer;
use ml::hyperopt::adapters::liquid::LiquidTrainer;
use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
use ml::hyperopt::adapters::tft::TFTTrainer;
use ml::hyperopt::adapters::tggn::TGGNTrainer;
use ml::hyperopt::adapters::tlob::TLOBTrainer;
use ml::hyperopt::adapters::xlstm::XLSTMTrainer;
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, liquid, tggn, tlob, kan, xlstm, diffusion, both, all)")]
struct Args {
    /// Model to optimize: "tft", "mamba2", "liquid", "tggn", "tlob", "kan",
    /// "xlstm", "diffusion", "both" (tft+mamba2), or "all" (all 8 models)
    #[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,
    })
}

#[allow(clippy::cognitive_complexity)]
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 mut 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);

    if let Err(e) = trainer.preload_data() {
        warn!("TFT data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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,
    ))
}

#[allow(clippy::cognitive_complexity)]
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 mut trainer = Mamba2Trainer::new(&args.data_dir, args.epochs)
        .context("Failed to create Mamba2 trainer")?
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("Mamba2 data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_liquid_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  Liquid Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-liquid");
    let training_paths = TrainingPaths::new(&args.base_dir, "liquid", &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 mut trainer = LiquidTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create Liquid trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("Liquid data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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("Liquid hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("Liquid 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,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_tggn_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  TGGN Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-tggn");
    let training_paths = TrainingPaths::new(&args.base_dir, "tggn", &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 mut trainer = TGGNTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create TGGN trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("TGGN data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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("TGGN hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("TGGN 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,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_tlob_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  TLOB Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-tlob");
    let training_paths = TrainingPaths::new(&args.base_dir, "tlob", &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 mut trainer = TLOBTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create TLOB trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("TLOB data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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("TLOB hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("TLOB 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,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_kan_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  KAN Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-kan");
    let training_paths = TrainingPaths::new(&args.base_dir, "kan", &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 mut trainer = KANTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create KAN trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("KAN data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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("KAN hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("KAN 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,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_xlstm_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  xLSTM Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-xlstm");
    let training_paths = TrainingPaths::new(&args.base_dir, "xlstm", &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 mut trainer = XLSTMTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create xLSTM trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("xLSTM data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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("xLSTM hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("xLSTM 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,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_diffusion_hyperopt(args: &Args) -> Result<Value> {
    info!("========================================");
    info!("  Diffusion Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-diffusion");
    let training_paths = TrainingPaths::new(&args.base_dir, "diffusion", &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 mut trainer = DiffusionTrainer::new(&args.data_dir, args.epochs)
        .context("Failed to create Diffusion trainer")?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    if let Err(e) = trainer.preload_data() {
        warn!("Diffusion data preload failed ({}), trials will load from disk", e);
    } else {
        info!("Training data preloaded for all {} trials", args.trials);
    }

    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("Diffusion hyperopt optimization failed")?;
    let elapsed = start.elapsed().as_secs_f64();

    info!("Diffusion 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,
    ))
}

const VALID_MODELS: &[&str] = &[
    "tft", "mamba2", "liquid", "tggn", "tlob", "kan", "xlstm", "diffusion",
];

#[allow(clippy::cognitive_complexity)]
fn main() -> Result<()> {
    // Initialize tracing with optional OTLP export to Tempo
    let otlp_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok();
    if let Err(e) = common::observability::init_observability(
        "hyperopt_baseline_supervised",
        otlp_endpoint.as_deref(),
    ) {
        eprintln!("Observability init failed (non-fatal): {e}");
    }

    training_metrics::init();
    metrics_server::start_metrics_server(9094);
    training_metrics::set_active_workers(1.0);

    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()
        );
    }

    // Resolve model selection: "both" = tft+mamba2 (backward compat), "all" = all 8 models
    let models: Vec<&str> = if args.model == "both" {
        vec!["tft", "mamba2"]
    } else if args.model == "all" {
        VALID_MODELS.to_vec()
    } else if VALID_MODELS.contains(&args.model.as_str()) {
        vec![args.model.as_str()]
    } else {
        anyhow::bail!(
            "Invalid --model value '{}'. Must be one of: {}, 'both', or 'all'.",
            args.model,
            VALID_MODELS.join(", ")
        );
    };

    info!("Models to optimize: {:?}", models);

    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();

    for model_name in &models {
        let result = match *model_name {
            "tft" => run_tft_hyperopt(&args),
            "mamba2" => run_mamba2_hyperopt(&args),
            "liquid" => run_liquid_hyperopt(&args),
            "tggn" => run_tggn_hyperopt(&args),
            "tlob" => run_tlob_hyperopt(&args),
            "kan" => run_kan_hyperopt(&args),
            "xlstm" => run_xlstm_hyperopt(&args),
            "diffusion" => run_diffusion_hyperopt(&args),
            other => {
                error!("Unknown model: {}", other);
                continue;
            }
        };

        match result {
            Ok(model_result) => {
                results.insert(model_name.to_string(), model_result);
            }
            Err(e) => {
                error!("{} hyperopt failed: {:#}", model_name, e);
                warn!("Continuing with remaining models...");
                results.insert(
                    model_name.to_string(),
                    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);

    training_metrics::set_active_workers(0.0);
    Ok(())
}