foxhunt/ml/src/hyperopt/campaign.rs

//! Hyperopt campaign configuration and runner.
//!
//! Defines campaign parameters for systematic hyperparameter search,
//! and provides a `run_campaign()` orchestrator that drives the optimizer
//! against the configured model adapter, persisting results to disk.

use std::path::PathBuf;

use serde::{Deserialize, Serialize};

use crate::gpu::capabilities::cached_capabilities;
use crate::gpu::memory_profile::{self, resolve_batch_size};

// Use canonical ModelType from crate root (re-exported from common)
pub use crate::ModelType;

/// Campaign configuration for multi-trial hyperparameter optimization.
#[derive(Debug, Clone)]
pub struct CampaignConfig {
    /// Which model to optimize.
    pub model_type: ModelType,
    /// Total number of trials to run.
    pub num_trials: usize,
    /// Path to training data directory.
    pub data_dir: PathBuf,
    /// Maximum batch size (bounded by GPU VRAM).
    pub max_batch_size: usize,
    /// Reduction factor η for SHA/Hyperband.
    pub early_stopping_eta: usize,
    /// Maximum epochs per trial.
    pub max_epochs_per_trial: usize,
    /// Base directory for results output.
    pub results_base_dir: PathBuf,
}

impl CampaignConfig {
    /// DQN campaign defaults (50 trials, SHA with η=3, GPU-adaptive batch size).
    pub fn dqn_default() -> Self {
        let caps = cached_capabilities();
        let max_batch = resolve_batch_size(caps, &memory_profile::estimates::DQN, 512);
        Self {
            model_type: ModelType::DQN,
            num_trials: 50,
            data_dir: PathBuf::from("test_data/real/databento/ml_training"),
            max_batch_size: max_batch,
            early_stopping_eta: 3,
            max_epochs_per_trial: 81,
            results_base_dir: PathBuf::from("ml/hyperopt_results"),
        }
    }

    /// PPO campaign defaults (30 trials, Hyperband, GPU-adaptive batch size).
    pub fn ppo_default() -> Self {
        let caps = cached_capabilities();
        let max_batch = resolve_batch_size(caps, &memory_profile::estimates::PPO, 512);
        Self {
            model_type: ModelType::PPO,
            num_trials: 30,
            data_dir: PathBuf::from("test_data/real/databento/ml_training"),
            max_batch_size: max_batch,
            early_stopping_eta: 3,
            max_epochs_per_trial: 81,
            results_base_dir: PathBuf::from("ml/hyperopt_results"),
        }
    }

    /// Generate timestamped results directory path.
    pub fn results_dir(&self) -> String {
        let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
        format!(
            "{}/{}/{}",
            self.results_base_dir.display(),
            self.model_type,
            timestamp
        )
    }
}

/// Results from a completed campaign.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CampaignResults {
    /// Number of trials that completed.
    pub trials_completed: usize,
    /// Best objective value found (lower is better).
    pub best_loss: f64,
    /// Best parameters serialized as JSON.
    pub best_params_json: String,
    /// Path to best checkpoint file, if saved.
    pub best_checkpoint_path: Option<String>,
    /// Total wall-clock time in seconds.
    pub total_time_seconds: f64,
}

/// Run a hyperopt campaign with the given configuration.
///
/// Creates an `ArgminOptimizer`, runs it against the configured model adapter,
/// and saves results to the campaign results directory.
///
/// # Errors
///
/// Returns an error if:
/// - The results directory cannot be created
/// - The model adapter fails to initialize
/// - The optimization run fails
/// - Result serialization or file writes fail
pub fn run_campaign(config: &CampaignConfig) -> anyhow::Result<CampaignResults> {
    let start = std::time::Instant::now();

    // Create results directory
    let results_dir = config.results_dir();
    std::fs::create_dir_all(&results_dir)?;

    match config.model_type {
        ModelType::DQN => run_dqn_campaign(config, &results_dir, start),
        ModelType::PPO => {
            anyhow::bail!("PPO campaign not yet implemented")
        }
        other => {
            anyhow::bail!("Hyperopt campaign not supported for model type: {}", other)
        }
    }
}

fn run_dqn_campaign(
    config: &CampaignConfig,
    results_dir: &str,
    start: std::time::Instant,
) -> anyhow::Result<CampaignResults> {
    use crate::hyperopt::adapters::dqn::DQNTrainer;
    use crate::hyperopt::optimizer::ArgminOptimizer;

    // Build optimizer
    let optimizer = ArgminOptimizer::builder()
        .max_trials(config.num_trials)
        .n_initial((config.num_trials / 2).max(1).min(5))
        .seed(42)
        .build();

    // Create DQN adapter
    let adapter = DQNTrainer::new(&config.data_dir, config.max_epochs_per_trial)?;

    // Run optimization
    let result = optimizer.optimize(adapter)?;

    // Serialize best params
    let best_params_json = serde_json::to_string_pretty(&result.best_params)?;

    // Save best parameters
    let params_path = format!("{}/best_params.json", results_dir);
    std::fs::write(&params_path, &best_params_json)?;

    let campaign_results = CampaignResults {
        trials_completed: result.all_trials.len(),
        best_loss: result.best_objective,
        best_params_json,
        best_checkpoint_path: None, // Checkpoints saved by trainer internally
        total_time_seconds: start.elapsed().as_secs_f64(),
    };

    // Save campaign summary
    let summary_path = format!("{}/campaign_summary.json", results_dir);
    std::fs::write(
        &summary_path,
        serde_json::to_string_pretty(&campaign_results)?,
    )?;

    Ok(campaign_results)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_campaign_config_dqn_defaults() {
        let config = CampaignConfig::dqn_default();
        assert_eq!(config.model_type, ModelType::DQN);
        assert_eq!(config.num_trials, 50);
        assert!(config.max_batch_size > 0, "max_batch_size should be positive");
        // No longer asserts <= 230 — batch size is dynamic based on GPU
    }

    #[test]
    fn test_campaign_config_ppo_defaults() {
        let config = CampaignConfig::ppo_default();
        assert_eq!(config.model_type, ModelType::PPO);
        assert_eq!(config.num_trials, 30);
    }

    #[test]
    fn test_results_dir_creation() {
        let config = CampaignConfig::dqn_default();
        let dir = config.results_dir();
        assert!(dir.starts_with("ml/hyperopt_results/dqn/"));
        // Should contain timestamp
        assert!(dir.len() > "ml/hyperopt_results/dqn/".len());
    }

    #[test]
    fn test_model_type_display() {
        assert_eq!(format!("{}", ModelType::DQN), "dqn");
        assert_eq!(format!("{}", ModelType::PPO), "ppo");
    }

    #[test]
    fn test_campaign_results_serialization() {
        let results = CampaignResults {
            trials_completed: 10,
            best_loss: 0.42,
            best_params_json: r#"{"learning_rate": 0.001}"#.to_string(),
            best_checkpoint_path: Some("path/to/best.safetensors".into()),
            total_time_seconds: 120.0,
        };
        let json = serde_json::to_string(&results).expect("serialize should succeed");
        let back: CampaignResults =
            serde_json::from_str(&json).expect("deserialize should succeed");
        assert_eq!(back.trials_completed, 10);
        assert!((back.best_loss - 0.42).abs() < 1e-10);
        assert_eq!(
            back.best_checkpoint_path,
            Some("path/to/best.safetensors".to_string())
        );
    }

    #[test]
    fn test_campaign_results_without_checkpoint() {
        let results = CampaignResults {
            trials_completed: 5,
            best_loss: 0.55,
            best_params_json: "{}".into(),
            best_checkpoint_path: None,
            total_time_seconds: 60.0,
        };
        let json = serde_json::to_string(&results).expect("serialize");
        assert!(json.contains("\"best_checkpoint_path\":null"));
    }
}