foxhunt/ml/examples/hyperopt_tft_demo.rs

//! TFT Hyperparameter Optimization Demo
//!
//! This example demonstrates how to use the argmin-based hyperparameter
//! optimization framework with Temporal Fusion Transformer (TFT). It runs
//! a small-scale optimization to show the complete workflow.
//!
//! ## Usage
//!
//! ```bash
//! # Local training with small dataset (default /tmp)
//! cargo run -p ml --example hyperopt_tft_demo --release --features cuda -- \
//!   --parquet-file test_data/ES_FUT_180d.parquet \
//!   --trials 10 \
//!   --epochs 20
//!
//! # Runpod training with custom base dir
//! ./hyperopt_tft_demo \
//!   --parquet-file /runpod-volume/datasets/parquet/futures/ES_FUT_180d.parquet \
//!   --base-dir /runpod-volume \
//!   --trials 50 \
//!   --epochs 50
//!
//! # Resume from specific run
//! ./hyperopt_tft_demo \
//!   --base-dir /runpod-volume \
//!   --run-id 20251028_223000_hyperopt \
//!   --parquet-file test_data/ES_FUT_180d.parquet \
//!   --trials 50 \
//!   --epochs 50
//! ```
//!
//! ## Output
//!
//! The example will:
//! 1. Initialize TFT trainer with specified Parquet file
//! 2. Run argmin optimization with Particle Swarm
//! 3. Display trial results including loss and parameter values
//! 4. Report best hyperparameters found
//! 5. Show expected improvement vs default parameters

use anyhow::Result;
use clap::Parser;
use ml::hyperopt::adapters::tft::TFTTrainer;
use ml::hyperopt::paths::{generate_run_id, TrainingPaths};
use ml::hyperopt::ArgminOptimizer;
use tracing::{info, Level};
use tracing_subscriber;

#[derive(Parser, Debug)]
#[command(name = "TFT Hyperparameter Optimization Demo")]
#[command(about = "Demonstrates argmin-based hyperparameter optimization for TFT")]
struct Args {
    /// Path to Parquet file with OHLCV data
    #[arg(long)]
    parquet_file: String,

    /// Number of optimization trials (default: 10)
    #[arg(long, default_value = "10")]
    trials: usize,

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

    /// Number of initial random samples (default: 3)
    #[arg(long, default_value = "3")]
    n_initial: usize,

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

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

    /// Maximum batch size for GPU memory constraints (default: 128 for RTX A4000 16GB)
    /// Examples: RTX 3050 Ti 4GB = 64, RTX A4000 16GB = 128, RTX 4090 24GB = 256
    #[arg(long, default_value = "128")]
    batch_size_max: usize,

    /// Base directory for training outputs (e.g., /runpod-volume)
    #[arg(long, default_value = "/tmp/ml_training")]
    base_dir: String,

    /// Run ID (auto-generated if not provided)
    #[arg(long)]
    run_id: Option<String>,

    /// Run type (hyperopt, production, test)
    #[arg(long, default_value = "hyperopt")]
    run_type: String,

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

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

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

    // Generate run ID if not provided
    let run_id = args
        .run_id
        .unwrap_or_else(|| generate_run_id(&args.run_type));

    // Create training paths
    let training_paths = TrainingPaths::new(&args.base_dir, "tft", &run_id);

    info!("========================================");
    info!("TFT Hyperparameter Optimization Demo");
    info!("========================================");
    info!("Configuration:");
    info!("  Parquet file: {}", args.parquet_file);
    info!("  Trials: {}", args.trials);
    info!("  Epochs per trial: {}", args.epochs);
    info!("  Initial samples: {}", args.n_initial);
    info!("  Random seed: {}", args.seed);
    info!(
        "  Batch size bounds: [{}, {}]",
        args.batch_size_min, args.batch_size_max
    );
    info!("");
    info!("Training Paths:");
    info!("  Run ID: {}", run_id);
    info!("  Base directory: {}", args.base_dir);
    info!("  Run directory: {:?}", training_paths.run_dir());
    info!("  Checkpoints: {:?}", training_paths.checkpoints_dir());
    info!("  Logs: {:?}", training_paths.logs_dir());
    info!("  Hyperopt: {:?}", training_paths.hyperopt_dir());
    info!("");

    // Create trainer with training paths
    info!("Creating TFT trainer...");
    let trainer = TFTTrainer::new(&args.parquet_file, args.epochs)?
        .with_early_stopping(args.early_stopping_patience)
        .with_training_paths(training_paths);

    info!("TFT Configuration:");
    info!("  Input features: 54 (Wave C + Wave D)");
    info!("  Sequence length: 60");
    info!("  Prediction horizon: 10");
    info!("  Quantiles: 3 (0.1, 0.5, 0.9)");
    info!("");

    // Create optimizer
    info!("Initializing argmin optimizer...");
    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    // Run optimization
    info!("");
    info!("Starting optimization (this may take a while)...");
    info!(
        "Expected runtime: ~{} minutes",
        estimate_runtime(args.trials, args.epochs)
    );
    info!("");

    let result = optimizer.optimize(trainer)?;

    // Display results
    info!("");
    info!("========================================");
    info!("Optimization Complete!");
    info!("========================================");
    info!("");
    info!("Best Hyperparameters:");
    info!("  Learning rate: {:.6}", result.best_params.learning_rate);
    info!("  Batch size: {}", result.best_params.batch_size);
    info!("  Hidden size: {}", result.best_params.hidden_size);
    info!("  Attention heads: {}", result.best_params.num_heads);
    info!("  Dropout: {:.3}", result.best_params.dropout);
    info!("");
    info!("Performance:");
    info!("  Best validation loss: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());

    // Find convergence trial (where best was found)
    let convergence_trial = result
        .all_trials
        .iter()
        .position(|t| (t.objective - result.best_objective).abs() < 1e-10)
        .unwrap_or(0);
    info!("  Convergence: {} trials to best", convergence_trial + 1);
    info!("");

    // Show top 5 trials
    if result.all_trials.len() >= 5 {
        info!("Top 5 Trials:");
        let mut sorted_trials = result.all_trials.clone();
        sorted_trials.sort_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap());

        for (i, trial) in sorted_trials.iter().take(5).enumerate() {
            info!(
                "  {}. Loss: {:.6} (LR: {:.6}, BS: {}, Hidden: {}, Heads: {})",
                i + 1,
                trial.objective,
                trial.params.learning_rate,
                trial.params.batch_size,
                trial.params.hidden_size,
                trial.params.num_heads
            );
        }
    }

    info!("");
    info!("========================================");
    info!("Architecture Insights:");
    info!("========================================");

    // Analyze best parameters
    let best = &result.best_params;

    // Calculate model complexity
    let complexity_score = (best.hidden_size as f64 * best.num_heads as f64) / 1000.0;
    let complexity_level = if complexity_score < 2.0 {
        "Light"
    } else if complexity_score < 4.0 {
        "Balanced"
    } else {
        "Heavy"
    };

    info!(
        "Model Complexity: {} (score: {:.2})",
        complexity_level, complexity_score
    );
    info!("  Hidden dimension: {} features", best.hidden_size);
    info!("  Attention heads: {} heads", best.num_heads);
    info!(
        "  Head dimension: {} features/head",
        best.hidden_size / best.num_heads
    );
    info!("");

    // Regularization analysis
    let regularization_level = if best.dropout < 0.1 {
        "Low"
    } else if best.dropout < 0.2 {
        "Medium"
    } else {
        "High"
    };

    info!("Regularization: {}", regularization_level);
    info!("  Dropout rate: {:.1}%", best.dropout * 100.0);
    info!("");

    // Training characteristics
    info!("Training Characteristics:");
    info!(
        "  Learning rate: {:.6} ({})",
        best.learning_rate,
        if best.learning_rate < 5e-5 {
            "Conservative"
        } else if best.learning_rate < 2e-4 {
            "Balanced"
        } else {
            "Aggressive"
        }
    );
    info!(
        "  Batch size: {} (GPU memory: ~{}MB)",
        best.batch_size,
        estimate_gpu_memory(best.batch_size, best.hidden_size)
    );
    info!("");

    info!("========================================");
    info!("Next Steps:");
    info!("========================================");
    info!("1. Use best parameters for production training");
    info!("2. Run longer optimization (50+ trials) for better results");
    info!("3. Validate on holdout dataset");
    info!("4. Deploy optimized model to trading system");
    info!(
        "5. Consider hidden_size={} as your production baseline",
        best.hidden_size
    );

    Ok(())
}

/// Estimate runtime based on trials and epochs
fn estimate_runtime(trials: usize, epochs: usize) -> usize {
    // Rough estimate: 2 min per 50 epochs on RTX 3050 Ti for TFT
    let minutes_per_trial = (epochs as f64 / 50.0) * 2.0;
    let total_minutes = (trials as f64 * minutes_per_trial).ceil() as usize;
    total_minutes
}

/// Estimate GPU memory usage for a given configuration
fn estimate_gpu_memory(batch_size: usize, hidden_size: usize) -> usize {
    // Rough estimate: base (200MB) + sequence memory
    // TFT has encoder-decoder architecture with attention
    let base_memory = 200;
    let sequence_memory = (batch_size * hidden_size * 60 * 8) / 1_000_000; // 60 seq length, 8 bytes/float
    let attention_memory = (batch_size * 60 * 60 * 4) / 1_000_000; // attention matrix

    base_memory + sequence_memory + attention_memory
}