foxhunt/ml/src/benchmarks.rs

//! ML Model Benchmark Suite
//!
//! Comprehensive benchmarking for all ML models with sub-50μs inference targets.
//! Validates GPU acceleration and performance requirements for HFT systems.

 // For canonical types

use std::time::Instant;

use anyhow::Result;
use candle_core::{Device, Tensor};
use tracing::{info, warn};

use crate::dqn::{RainbowAgent, RainbowAgentConfig};
use crate::liquid::{LiquidNetwork, LiquidNetworkConfig, NetworkType};
use crate::mamba::{Mamba2Config, Mamba2SSM};
use crate::ppo::{ContinuousPPO, ContinuousPPOConfig};
use crate::tft::{TFTConfig, TemporalFusionTransformer};
use crate::tlob::{TLOBConfig, TLOBTransformer};
use crate::MLError;

/// Benchmark configuration
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
    pub warmup_runs: usize,
    pub test_runs: usize,
    pub batch_size: usize,
    pub target_latency_us: u64,
    pub enable_gpu: bool,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            warmup_runs: 100,
            test_runs: 1000,
            batch_size: 1,
            target_latency_us: 50,
            enable_gpu: true,
        }
    }
}

/// Benchmark results for a single model
#[derive(Debug, Clone)]
pub struct ModelBenchmarkResults {
    pub model_name: String,
    pub device: String,
    pub avg_latency_us: f64,
    pub p50_latency_us: f64,
    pub p95_latency_us: f64,
    pub p99_latency_us: f64,
    pub max_latency_us: f64,
    pub min_latency_us: f64,
    pub throughput_pps: f64,
    pub target_met: bool,
    pub compilation_time_ms: f64,
    pub memory_usage_mb: f64,
    pub gpu_utilization_percent: f64,
}

/// Complete benchmark suite results
#[derive(Debug, Clone)]
pub struct BenchmarkSuite {
    pub results: Vec<ModelBenchmarkResults>,
    pub system_info: SystemInfo,
    pub gpu_info: Option<GpuInfo>,
    pub total_duration_ms: f64,
}

#[derive(Debug, Clone)]
pub struct SystemInfo {
    pub cpu_count: usize,
    pub available_memory_gb: f64,
    pub os: String,
    pub architecture: String,
}

#[derive(Debug, Clone)]
pub struct GpuInfo {
    pub name: String,
    pub memory_gb: f64,
    pub compute_capability: String,
    pub driver_version: String,
}

/// Main benchmark runner
pub struct MLBenchmarkRunner {
    config: BenchmarkConfig,
    device: Device,
}

impl MLBenchmarkRunner {
    pub fn new(config: BenchmarkConfig) -> Result<Self, MLError> {
        let device = if config.enable_gpu {
            Device::cuda_if_available(0).unwrap_or(Device::Cpu)
        } else {
            Device::Cpu
        };

        info!("Initialized ML Benchmark Runner on device: {:?}", device);

        Ok(Self { config, device })
    }

    /// Run complete benchmark suite
    pub async fn run_full_benchmark(&self) -> Result<BenchmarkSuite, MLError> {
        info!("Starting ML model benchmark suite");
        let start_time = Instant::now();

        let mut results = Vec::new();

        // Benchmark MAMBA-2 SSM
        if let Ok(mamba_results) = self.benchmark_mamba2().await {
            results.push(mamba_results);
        }

        // Benchmark DQN (Rainbow)
        if let Ok(dqn_results) = self.benchmark_dqn().await {
            results.push(dqn_results);
        }

        // Benchmark PPO
        if let Ok(ppo_results) = self.benchmark_ppo().await {
            results.push(ppo_results);
        }

        // Benchmark TLOB Transformer
        if let Ok(tlob_results) = self.benchmark_tlob().await {
            results.push(tlob_results);
        }

        // Benchmark TFT
        if let Ok(tft_results) = self.benchmark_tft().await {
            results.push(tft_results);
        }

        // Benchmark Liquid Networks
        if let Ok(liquid_results) = self.benchmark_liquid().await {
            results.push(liquid_results);
        }

        let total_duration = start_time.elapsed().as_millis() as f64;

        let suite = BenchmarkSuite {
            results,
            system_info: self.get_system_info(),
            gpu_info: self.get_gpu_info(),
            total_duration_ms: total_duration,
        };

        info!(
            "Benchmark suite completed in {:.2}ms with {} models",
            total_duration,
            suite.results.len()
        );

        Ok(suite)
    }

    /// Benchmark MAMBA-2 SSM
    pub async fn benchmark_mamba2(&self) -> Result<ModelBenchmarkResults, MLError> {
        info!("Benchmarking MAMBA-2 SSM");

        let config = Mamba2Config {
            d_model: 256,
            d_state: 32,
            num_layers: 4,
            target_latency_us: self.config.target_latency_us,
            batch_size: self.config.batch_size,
            hardware_aware: true,
            use_ssd: true,
            use_selective_state: true,
            ..Default::default()
        };

        let compilation_start = Instant::now();
        let mut model = Mamba2SSM::new(config)?;
        let compilation_time = compilation_start.elapsed().as_millis() as f64;

        // Generate test data
        let input_data: Vec<f64> = (0..256).map(|i| (i as f64) / 256.0).collect();

        // Warmup
        for _ in 0..self.config.warmup_runs {
            let _ = model.predict_single_fast(&input_data)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..self.config.test_runs {
            let start = Instant::now();
            let _ = model.predict_single_fast(&input_data)?;
            latencies.push(start.elapsed().as_micros() as f64);
        }

        Ok(self.calculate_results("MAMBA-2 SSM", latencies, compilation_time))
    }

    /// Benchmark DQN (Rainbow)
    pub async fn benchmark_dqn(&self) -> Result<ModelBenchmarkResults, MLError> {
        info!("Benchmarking Rainbow DQN");

        let mut config = RainbowAgentConfig::default();
        config.network_config.input_size = 64;
        config.network_config.num_actions = 4;
        config.learning_rate = 1e-4;
        config.batch_size = self.config.batch_size;
        config.device = if matches!(self.device, Device::Cuda(_)) {
            "cuda".to_string()
        } else {
            "cpu".to_string()
        };

        let compilation_start = Instant::now();
        let agent = RainbowAgent::new(config)?;
        let compilation_time = compilation_start.elapsed().as_millis() as f64;

        // Generate test data
        let state: Vec<f32> = (0..64).map(|i| (i as f32) / 64.0).collect();

        // Warmup
        for _ in 0..self.config.warmup_runs {
            let _ = agent.select_action(&state)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..self.config.test_runs {
            let start = Instant::now();
            let _ = agent.select_action(&state)?;
            latencies.push(start.elapsed().as_micros() as f64);
        }

        Ok(self.calculate_results("Rainbow DQN", latencies, compilation_time))
    }

    /// Benchmark PPO
    pub async fn benchmark_ppo(&self) -> Result<ModelBenchmarkResults, MLError> {
        info!("Benchmarking PPO");

        let mut config = ContinuousPPOConfig::default();
        config.state_dim = 32;
        config.policy_learning_rate = 1e-4;
        config.value_learning_rate = 1e-4;
        config.batch_size = self.config.batch_size;

        let compilation_start = Instant::now();
        let ppo = ContinuousPPO::new(config)?;
        let compilation_time = compilation_start.elapsed().as_millis() as f64;

        // Generate test data
        let state = vec![0.1_f32; 32]; // PPO expects &[f32]

        // Warmup
        for _ in 0..self.config.warmup_runs {
            let _ = ppo.act(&state)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..self.config.test_runs {
            let start = Instant::now();
            let _ = ppo.act(&state)?;
            latencies.push(start.elapsed().as_micros() as f64);
        }

        Ok(self.calculate_results("PPO", latencies, compilation_time))
    }

    /// Benchmark TLOB Transformer
    pub async fn benchmark_tlob(&self) -> Result<ModelBenchmarkResults, MLError> {
        info!("Benchmarking TLOB Transformer");

        let mut config = TLOBConfig::default();
        config.feature_dim = 20;
        config.batch_size = self.config.batch_size;

        let compilation_start = Instant::now();
        let tlob = TLOBTransformer::new(config)?;
        let compilation_time = compilation_start.elapsed().as_millis() as f64;

        // Generate test order book data - using transformer TLOBFeatures
        let tlob_features = crate::tlob::transformer::TLOBFeatures {
            timestamp: 1642531200000,
            bid_prices: [100000; 10],
            ask_prices: [100010; 10],
            bid_sizes: [1000; 10],
            ask_sizes: [1000; 10],
            trade_price: 100005,
            trade_size: 500,
            spread: 10,
            mid_price: 100005,
            microstructure_features: [1, 2, 3],
        };

        // Warmup
        for _ in 0..self.config.warmup_runs {
            let _ = tlob.predict(&tlob_features)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..self.config.test_runs {
            let start = Instant::now();
            let _ = tlob.predict(&tlob_features)?;
            latencies.push(start.elapsed().as_micros() as f64);
        }

        Ok(self.calculate_results("TLOB Transformer", latencies, compilation_time))
    }

    /// Benchmark TFT
    pub async fn benchmark_tft(&self) -> Result<ModelBenchmarkResults, MLError> {
        info!("Benchmarking Temporal Fusion Transformer");

        let config = TFTConfig {
            input_dim: 64,
            hidden_dim: 128,
            num_heads: 8,
            num_layers: 3,
            prediction_horizon: 10,
            sequence_length: 50,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 20,
            batch_size: self.config.batch_size,
            max_inference_latency_us: self.config.target_latency_us,
            ..Default::default()
        };

        let compilation_start = Instant::now();
        let mut tft = TemporalFusionTransformer::new(config)?;
        let compilation_time = compilation_start.elapsed().as_millis() as f64;

        // Generate test data
        let static_features: Vec<f32> = (0..5).map(|i| i as f32 / 5.0).collect();
        let historical_features: Vec<f32> = (0..1000).map(|i| i as f32 / 1000.0).collect(); // 50x20
        let future_features: Vec<f32> = (0..100).map(|i| i as f32 / 100.0).collect(); // 10x10

        // Warmup
        for _ in 0..self.config.warmup_runs {
            let _ = tft.predict_fast(&static_features, &historical_features, &future_features)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..self.config.test_runs {
            let start = Instant::now();
            let _ = tft.predict_fast(&static_features, &historical_features, &future_features)?;
            latencies.push(start.elapsed().as_micros() as f64);
        }

        Ok(self.calculate_results("TFT", latencies, compilation_time))
    }

    /// Benchmark Liquid Networks
    pub async fn benchmark_liquid(&self) -> Result<ModelBenchmarkResults, MLError> {
        info!("Benchmarking Liquid Neural Networks");

        use crate::liquid::{ActivationType, FixedPoint, LTCConfig, SolverType, PRECISION};

        let ltc_config = LTCConfig {
            input_size: 32,
            hidden_size: 64,
            tau_min: FixedPoint(PRECISION / 100),
            tau_max: FixedPoint(PRECISION),
            use_bias: true,
            solver_type: SolverType::Euler,
            activation: ActivationType::Tanh,
        };

        let config = LiquidNetworkConfig {
            network_type: NetworkType::CfC,
            input_size: 32,
            output_size: 4,
            layer_configs: vec![crate::liquid::LayerConfig::LTC(ltc_config)],
            output_layer: crate::liquid::OutputLayerConfig {
                use_linear_output: true,
                output_activation: None,
                dropout_rate: None,
            },
            default_dt: FixedPoint(PRECISION / 100),
            market_regime_adaptation: true,
        };

        let compilation_start = Instant::now();
        let mut liquid = LiquidNetwork::new(config)?;
        let compilation_time = compilation_start.elapsed().as_millis() as f64;

        // Generate test data
        let input_data: Vec<f64> = (0..32).map(|i| (i as f64) / 32.0).collect();

        // Warmup
        for _ in 0..self.config.warmup_runs {
            let _ = liquid.predict(&input_data)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..self.config.test_runs {
            let start = Instant::now();
            let _ = liquid.predict(&input_data)?;
            latencies.push(start.elapsed().as_micros() as f64);
        }

        Ok(self.calculate_results("Liquid Networks", latencies, compilation_time))
    }

    fn calculate_results(
        &self,
        model_name: &str,
        mut latencies: Vec<f64>,
        compilation_time_ms: f64,
    ) -> ModelBenchmarkResults {
        latencies.sort_by(|a, b| a.partial_cmp(b).unwrap());

        let avg_latency = latencies.iter().sum::<f64>() / latencies.len() as f64;
        let min_latency = latencies[0];
        let max_latency = latencies[latencies.len() - 1];

        let p50_idx = latencies.len() / 2;
        let p95_idx = (latencies.len() * 95) / 100;
        let p99_idx = (latencies.len() * 99) / 100;

        let p50_latency = latencies[p50_idx];
        let p95_latency = latencies[p95_idx.min(latencies.len() - 1)];
        let p99_latency = latencies[p99_idx.min(latencies.len() - 1)];

        let throughput = if avg_latency > 0.0 {
            1_000_000.0 / avg_latency // predictions per second
        } else {
            0.0
        };

        let target_met = avg_latency <= self.config.target_latency_us as f64;

        let device_name = match &self.device {
            Device::Cpu => "CPU",
            Device::Cuda(_) => "CUDA",
            Device::Metal(_) => "Metal",
        }
        .to_string();

        if !target_met {
            warn!(
                "{} average latency {:.2}μs exceeds target {}μs",
                model_name, avg_latency, self.config.target_latency_us
            );
        } else {
            info!(
                "{} meets target: {:.2}μs average latency, {:.0} pps throughput",
                model_name, avg_latency, throughput
            );
        }

        ModelBenchmarkResults {
            model_name: model_name.to_string(),
            device: device_name,
            avg_latency_us: avg_latency,
            p50_latency_us: p50_latency,
            p95_latency_us: p95_latency,
            p99_latency_us: p99_latency,
            max_latency_us: max_latency,
            min_latency_us: min_latency,
            throughput_pps: throughput,
            target_met,
            compilation_time_ms,
            memory_usage_mb: 0.0,         // TODO: Implement memory measurement
            gpu_utilization_percent: 0.0, // TODO: Implement GPU utilization measurement
        }
    }

    fn get_system_info(&self) -> SystemInfo {
        SystemInfo {
            cpu_count: num_cpus::get(),
            available_memory_gb: 8.0, // TODO: Get actual memory info
            os: std::env::consts::OS.to_string(),
            architecture: std::env::consts::ARCH.to_string(),
        }
    }

    fn get_gpu_info(&self) -> Option<GpuInfo> {
        if matches!(self.device, Device::Cuda(_)) {
            Some(GpuInfo {
                name: "CUDA Device".to_string(),       // TODO: Get actual GPU name
                memory_gb: 8.0,                        // TODO: Get actual GPU memory
                compute_capability: "8.0".to_string(), // TODO: Get actual compute capability
                driver_version: "Unknown".to_string(), // TODO: Get actual driver version
            })
        } else {
            None
        }
    }

    /// Generate performance report
    pub fn generate_report(&self, suite: &BenchmarkSuite) -> String {
        let mut report = String::new();
        report.push_str("# ML Model Performance Benchmark Report\n\n");

        // System Information
        report.push_str("## System Information\n");
        report.push_str(&format!("- OS: {}\n", suite.system_info.os));
        report.push_str(&format!(
            "- Architecture: {}\n",
            suite.system_info.architecture
        ));
        report.push_str(&format!("- CPU Cores: {}\n", suite.system_info.cpu_count));
        report.push_str(&format!(
            "- Available Memory: {:.1} GB\n",
            suite.system_info.available_memory_gb
        ));

        if let Some(gpu) = &suite.gpu_info {
            report.push_str(&format!("- GPU: {}\n", gpu.name));
            report.push_str(&format!("- GPU Memory: {:.1} GB\n", gpu.memory_gb));
            report.push_str(&format!(
                "- Compute Capability: {}\n",
                gpu.compute_capability
            ));
        }

        report.push_str(&format!("\n## Benchmark Configuration\n"));
        report.push_str(&format!(
            "- Target Latency: {}μs\n",
            self.config.target_latency_us
        ));
        report.push_str(&format!("- Test Runs: {}\n", self.config.test_runs));
        report.push_str(&format!("- Warmup Runs: {}\n", self.config.warmup_runs));
        report.push_str(&format!("- Batch Size: {}\n", self.config.batch_size));

        report.push_str("\n## Performance Results\n\n");
        report.push_str("| Model | Device | Avg (μs) | P95 (μs) | P99 (μs) | Max (μs) | Throughput (pps) | Target Met |\n");
        report.push_str("|-------|--------|----------|----------|----------|----------|------------------|------------|\n");

        for result in &suite.results {
            let target_met = if result.target_met { "✅" } else { "❌" };
            report.push_str(&format!(
                "| {} | {} | {:.1} | {:.1} | {:.1} | {:.1} | {:.0} | {} |\n",
                result.model_name,
                result.device,
                result.avg_latency_us,
                result.p95_latency_us,
                result.p99_latency_us,
                result.max_latency_us,
                result.throughput_pps,
                target_met
            ));
        }

        let models_meeting_target = suite.results.iter().filter(|r| r.target_met).count();
        report.push_str(&format!(
            "\n**Summary**: {}/{} models meet the <{}μs latency target\n",
            models_meeting_target,
            suite.results.len(),
            self.config.target_latency_us
        ));

        report.push_str(&format!(
            "\nTotal benchmark time: {:.2}ms\n",
            suite.total_duration_ms
        ));

        report
    }
}

/// Quick GPU capability test
pub fn test_gpu_acceleration() -> Result<bool, MLError> {
    info!("Testing GPU acceleration capabilities");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    match device {
        Device::Cuda(_) => {
            info!("CUDA GPU detected and available");

            // Test basic tensor operations on GPU
            let a = Tensor::randn(0.0, 1.0, (1000, 1000), &device).map_err(|e| {
                MLError::TensorCreationError {
                    operation: "GPU test tensor A".to_string(),
                    reason: e.to_string(),
                }
            })?;
            let b = Tensor::randn(0.0, 1.0, (1000, 1000), &device).map_err(|e| {
                MLError::TensorCreationError {
                    operation: "GPU test tensor B".to_string(),
                    reason: e.to_string(),
                }
            })?;

            let start = Instant::now();
            let _ = a
                .matmul(&b)
                .map_err(|e| MLError::ModelError(format!("GPU matmul test failed: {}", e)))?;
            let gpu_time = start.elapsed();

            info!("GPU matrix multiplication test completed in {:?}", gpu_time);
            Ok(true)
        }
        _ => {
            info!("No CUDA GPU available, using CPU");
            Ok(false)
        }
    }
}

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

    #[tokio::test]
    async fn test_benchmark_runner_creation() {
        let config = BenchmarkConfig::default();
        let runner = MLBenchmarkRunner::new(config).unwrap();
        assert!(runner.config.test_runs > 0);
    }

    #[tokio::test]
    async fn test_gpu_detection() {
        let result = test_gpu_acceleration();
        assert!(result.is_ok());
    }

    #[test]
    fn test_benchmark_config_default() {
        let config = BenchmarkConfig::default();
        assert_eq!(config.target_latency_us, 50);
        assert!(config.test_runs > 0);
        assert!(config.warmup_runs > 0);
    }
}