//! 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;
use crate::batch_size_resolver::resolve_batch_size;

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

/// Campaign configuration for multi-trial hyperparameter optimization.
/// Campaign execution mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CampaignMode {
    /// Quick: fewer trials, fewer epochs. For local validation.
    Quick,
    /// Standard: default trials and epochs. For single-phase hyperopt.
    Standard,
    /// Full: all three training phases (BC → Online RL → Refinement).
    /// Runs max_epochs=100 per trial to cover all phases.
    /// Phase 1 (BC): epochs 0..c51_warmup_epochs (MSE + expert demos + DT)
    /// Phase 2 (RL): epochs warmup..80% (all features, C51 ramp, HER)
    /// Phase 3 (Refine): last 20% (pure C51, shrink-and-perturb)
    Full,
}

#[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,
    /// Campaign execution mode.
    pub mode: CampaignMode,
    /// Optional MBP-10 order book data directory for OFI features.
    /// When set, enables 8 additional OFI features (VPIN, Kyle's Lambda, etc.)
    /// producing state_dim=80 instead of 72.
    pub mbp10_data_dir: Option<PathBuf>,
    /// Optional trades data directory for VPIN/Kyle's Lambda computation.
    /// Required alongside mbp10_data_dir for full OFI feature extraction.
    pub trades_data_dir: Option<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"),
            mode: CampaignMode::Standard,
            mbp10_data_dir: None,
            trades_data_dir: None,
        }
    }

    /// DQN full pipeline: all 3 training phases (BC → RL → Refinement).
    /// 20 trials × 50 epochs — covers Phase 1 (warmup), Phase 2 (full RL),
    /// and Phase 3 (refinement with shrink-and-perturb).
    pub fn dqn_full() -> Self {
        let caps = cached_capabilities();
        let max_batch = resolve_batch_size(caps, &memory_profile::estimates::DQN, 512);
        Self {
            model_type: ModelType::DQN,
            num_trials: 20,
            data_dir: PathBuf::from("test_data/real/databento/ml_training"),
            max_batch_size: max_batch,
            early_stopping_eta: 3,
            max_epochs_per_trial: 50,
            results_base_dir: PathBuf::from("ml/hyperopt_results"),
            mode: CampaignMode::Full,
            mbp10_data_dir: None,
            trades_data_dir: None,
        }
    }

    /// DQN local dev campaign (10 trials × 20 epochs — RTX 3050 friendly).
    pub fn dqn_localdev() -> Self {
        Self {
            model_type: ModelType::DQN,
            num_trials: 10,
            data_dir: PathBuf::from("test_data/futures-baseline"),
            max_batch_size: 64,
            early_stopping_eta: 3,
            max_epochs_per_trial: 20,
            results_base_dir: PathBuf::from("ml/hyperopt_results"),
            mode: CampaignMode::Standard,
            mbp10_data_dir: None,
            trades_data_dir: None,
        }
    }

    /// 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"),
            mode: CampaignMode::Standard,
            mbp10_data_dir: None,
            trades_data_dir: None,
        }
    }

    /// 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 @ ModelType::CompactDQN
        | other @ ModelType::DistilledMicroNet
        | other @ ModelType::RainbowDQN
        | other @ ModelType::MAMBA
        | other @ ModelType::TFT
        | other @ ModelType::TGGN
        | other @ ModelType::LNN
        | other @ ModelType::TLOB
        | other @ ModelType::Transformer
        | other @ ModelType::Mamba
        | other @ ModelType::LiquidNet
        | other @ ModelType::TGNN
        | other @ ModelType::Ensemble
        | other @ ModelType::KAN
        | other @ ModelType::XLSTM
        | other @ ModelType::Diffusion => {
            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 with optional MBP-10 + trades data for OFI features
    let adapter = DQNTrainer::new(&config.data_dir, config.max_epochs_per_trial)?
        .with_ofi_data_dirs(
            config.mbp10_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
            config.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
        );

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

    #[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"));
    }

    /// Local hyperopt: 5 PSO trials × 50 epochs on real ES data.
    /// Run with: FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib --release -- test_local_hyperopt --ignored --nocapture
    #[test]
    #[ignore] // GPU + real data, ~30 min on RTX 3050
    fn test_local_hyperopt() {
        // Initialize tracing so info!/warn!/error! output is visible in test runs
        let _ = tracing_subscriber::fmt()
            .with_env_filter("info")
            .with_test_writer()
            .try_init();

        let data_dir = std::env::var("FOXHUNT_TEST_DATA")
            .unwrap_or_else(|_| "test_data/futures-baseline".to_owned());
        let data_path = std::path::Path::new(&data_dir);

        // Auto-detect: if data_dir has ES.FUT subdirectory, use it directly
        let effective_dir = if data_path.join("ES.FUT").exists() {
            data_path.join("ES.FUT")
        } else {
            data_path.to_path_buf()
        };

        if !effective_dir.exists() {
            eprintln!("Skipping: data dir not found: {}", effective_dir.display());
            return;
        }

        let mut config = CampaignConfig::dqn_localdev();
        config.data_dir = effective_dir;
        // Env var overrides for CI (quick) vs manual (full) runs
        if let Ok(t) = std::env::var("FOXHUNT_HYPEROPT_TRIALS") {
            config.num_trials = t.parse().unwrap_or(config.num_trials);
        }
        if let Ok(e) = std::env::var("FOXHUNT_HYPEROPT_EPOCHS") {
            config.max_epochs_per_trial = e.parse().unwrap_or(config.max_epochs_per_trial);
        }

        println!("Starting local DQN hyperopt: {} trials × {} epochs ({:?} mode)",
            config.num_trials, config.max_epochs_per_trial, config.mode);
        println!("Data: {}", config.data_dir.display());

        let result = run_campaign(&config).expect("campaign should complete");

        println!("\n=== HYPEROPT RESULTS ===");
        println!("Trials completed: {}", result.trials_completed);
        println!("Best loss (negative Sharpe): {:.4}", result.best_loss);
        println!("Total time: {:.1}s", result.total_time_seconds);
        println!("Best params:\n{}", result.best_params_json);

        assert!(result.trials_completed >= 1, "Should complete at least 1 trial");
    }
}