//! Memory profiling tool for ML models
//!
//! Measures memory usage, identifies hotspots, and validates optimizations.

use candle_core::{Device, DType, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Instant;
use sysinfo::{System, SystemExt, ProcessExt};

// Import model types
use ml::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::ppo::{WorkingPPO, PPOConfig};
use ml::tft::{TemporalFusionTransformer, TFTConfig};
use ml::mamba::{Mamba2SSM, Mamba2Config};
use ml::liquid::network::{LiquidNetwork, LiquidNetworkConfig};

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ModelMemoryProfile {
    model_name: String,
    base_memory_mb: f64,
    peak_memory_mb: f64,
    weight_memory_mb: f64,
    activation_memory_mb: f64,
    parameter_count: usize,
    inference_latency_us: u64,
    memory_per_parameter_bytes: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct MemoryOptimizationReport {
    timestamp: String,
    baseline_profiles: Vec<ModelMemoryProfile>,
    optimized_profiles: Vec<ModelMemoryProfile>,
    memory_savings_mb: HashMap<String, f64>,
    memory_reduction_percent: HashMap<String, f64>,
    accuracy_impact_percent: HashMap<String, f64>,
    recommendations: Vec<String>,
}

/// Measure current process memory usage
fn measure_memory_mb(sys: &mut System) -> f64 {
    sys.refresh_process(sysinfo::get_current_pid().unwrap());
    if let Some(process) = sys.process(sysinfo::get_current_pid().unwrap()) {
        process.memory() as f64 / 1_048_576.0 // Convert bytes to MB
    } else {
        0.0
    }
}

/// Profile DQN model memory
fn profile_dqn(device: &Device) -> Result<ModelMemoryProfile, Box<dyn std::error::Error>> {
    let mut sys = System::new_all();

    // Measure baseline memory
    let baseline_mb = measure_memory_mb(&mut sys);

    // Create DQN model
    let config = WorkingDQNConfig {
        state_dim: 256,
        action_dim: 3,
        hidden_dims: vec![512, 512, 256],
        learning_rate: 0.0003,
        gamma: 0.99,
        target_update_freq: 1000,
        batch_size: 64,
        buffer_capacity: 100_000,
    };

    let model = WorkingDQN::new(config, device.clone())?;

    // Measure model memory
    std::thread::sleep(std::time::Duration::from_millis(100));
    let model_mb = measure_memory_mb(&mut sys);
    let weight_memory_mb = model_mb - baseline_mb;

    // Measure inference memory (peak)
    let input = Tensor::randn(0f32, 1f32, (1, 256), device)?;
    let start = Instant::now();
    let _ = model.predict_action(&input)?;
    let inference_latency_us = start.elapsed().as_micros() as u64;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let peak_mb = measure_memory_mb(&mut sys);
    let activation_memory_mb = peak_mb - model_mb;

    // Count parameters
    let parameter_count = model.parameter_count();

    Ok(ModelMemoryProfile {
        model_name: "DQN".to_string(),
        base_memory_mb: baseline_mb,
        peak_memory_mb: peak_mb,
        weight_memory_mb,
        activation_memory_mb,
        parameter_count,
        inference_latency_us,
        memory_per_parameter_bytes: (weight_memory_mb * 1_048_576.0) / parameter_count as f64,
    })
}

/// Profile PPO model memory
fn profile_ppo(device: &Device) -> Result<ModelMemoryProfile, Box<dyn std::error::Error>> {
    let mut sys = System::new_all();

    let baseline_mb = measure_memory_mb(&mut sys);

    let config = PPOConfig {
        state_dim: 256,
        action_dim: 3,
        actor_hidden_dims: vec![512, 512, 256],
        critic_hidden_dims: vec![512, 512, 256],
        learning_rate: 0.0003,
        gamma: 0.99,
        gae_lambda: 0.95,
        clip_epsilon: 0.2,
        value_loss_coef: 0.5,
        entropy_coef: 0.01,
        max_grad_norm: 0.5,
        batch_size: 64,
        num_epochs: 10,
    };

    let model = WorkingPPO::new(config, device.clone())?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let model_mb = measure_memory_mb(&mut sys);
    let weight_memory_mb = model_mb - baseline_mb;

    let state = Tensor::randn(0f32, 1f32, (1, 256), device)?;
    let start = Instant::now();
    let _ = model.select_action(&state)?;
    let inference_latency_us = start.elapsed().as_micros() as u64;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let peak_mb = measure_memory_mb(&mut sys);
    let activation_memory_mb = peak_mb - model_mb;

    let parameter_count = model.parameter_count();

    Ok(ModelMemoryProfile {
        model_name: "PPO".to_string(),
        base_memory_mb: baseline_mb,
        peak_memory_mb: peak_mb,
        weight_memory_mb,
        activation_memory_mb,
        parameter_count,
        inference_latency_us,
        memory_per_parameter_bytes: (weight_memory_mb * 1_048_576.0) / parameter_count as f64,
    })
}

/// Profile TFT model memory
fn profile_tft(device: &Device) -> Result<ModelMemoryProfile, Box<dyn std::error::Error>> {
    let mut sys = System::new_all();

    let baseline_mb = measure_memory_mb(&mut sys);

    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 256,
        num_heads: 8,
        num_layers: 4,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 20,
        learning_rate: 1e-3,
        batch_size: 32,
        dropout_rate: 0.1,
        l2_regularization: 1e-4,
        use_flash_attention: true,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 50,
        target_throughput_pps: 100_000,
    };

    let mut model = TemporalFusionTransformer::new(config.clone())?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let model_mb = measure_memory_mb(&mut sys);
    let weight_memory_mb = model_mb - baseline_mb;

    // Create dummy inputs
    let static_features = vec![0.0f32; config.num_static_features];
    let historical_features = vec![0.0f32; config.sequence_length * config.num_unknown_features];
    let future_features = vec![0.0f32; config.prediction_horizon * config.num_known_features];

    let start = Instant::now();
    let _ = model.predict_fast(&static_features, &historical_features, &future_features)?;
    let inference_latency_us = start.elapsed().as_micros() as u64;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let peak_mb = measure_memory_mb(&mut sys);
    let activation_memory_mb = peak_mb - model_mb;

    // Estimate parameter count
    let parameter_count = estimate_tft_parameters(&config);

    Ok(ModelMemoryProfile {
        model_name: "TFT".to_string(),
        base_memory_mb: baseline_mb,
        peak_memory_mb: peak_mb,
        weight_memory_mb,
        activation_memory_mb,
        parameter_count,
        inference_latency_us,
        memory_per_parameter_bytes: (weight_memory_mb * 1_048_576.0) / parameter_count as f64,
    })
}

/// Profile MAMBA-2 model memory
fn profile_mamba(device: &Device) -> Result<ModelMemoryProfile, Box<dyn std::error::Error>> {
    let mut sys = System::new_all();

    let baseline_mb = measure_memory_mb(&mut sys);

    let config = Mamba2Config {
        d_model: 256,
        d_state: 64,
        d_head: 32,
        num_heads: 8,
        expand: 2,
        num_layers: 4,
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: true,
        target_latency_us: 5,
        max_seq_len: 512,
        learning_rate: 1e-3,
        weight_decay: 1e-4,
        grad_clip: 1.0,
        warmup_steps: 1000,
        batch_size: 16,
        seq_len: 256,
    };

    let mut model = Mamba2SSM::new(config.clone(), device)?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let model_mb = measure_memory_mb(&mut sys);
    let weight_memory_mb = model_mb - baseline_mb;

    let input = vec![0.0; config.d_model];
    let start = Instant::now();
    let _ = model.predict_single_fast(&input)?;
    let inference_latency_us = start.elapsed().as_micros() as u64;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let peak_mb = measure_memory_mb(&mut sys);
    let activation_memory_mb = peak_mb - model_mb;

    let parameter_count = model.metadata.num_parameters;

    Ok(ModelMemoryProfile {
        model_name: "MAMBA-2".to_string(),
        base_memory_mb: baseline_mb,
        peak_memory_mb: peak_mb,
        weight_memory_mb,
        activation_memory_mb,
        parameter_count,
        inference_latency_us,
        memory_per_parameter_bytes: (weight_memory_mb * 1_048_576.0) / parameter_count as f64,
    })
}

/// Profile Liquid model memory
fn profile_liquid() -> Result<ModelMemoryProfile, Box<dyn std::error::Error>> {
    let mut sys = System::new_all();

    let baseline_mb = measure_memory_mb(&mut sys);

    let config = LiquidNetworkConfig {
        input_dim: 256,
        hidden_dim: 512,
        output_dim: 3,
        num_layers: 4,
        cell_type: ml::liquid::cells::CellType::LTC,
        activation: ml::liquid::activation::ActivationType::Tanh,
        solver: ml::liquid::ode_solvers::SolverType::Euler,
        time_step: 0.001,
        inference_steps: 10,
    };

    let model = LiquidNetwork::new(config.clone())?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let model_mb = measure_memory_mb(&mut sys);
    let weight_memory_mb = model_mb - baseline_mb;

    let input = vec![ml::liquid::FixedPoint::zero(); config.input_dim];
    let start = Instant::now();
    let _ = model.forward(&input)?;
    let inference_latency_us = start.elapsed().as_micros() as u64;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let peak_mb = measure_memory_mb(&mut sys);
    let activation_memory_mb = peak_mb - model_mb;

    let parameter_count = estimate_liquid_parameters(&config);

    Ok(ModelMemoryProfile {
        model_name: "Liquid".to_string(),
        base_memory_mb: baseline_mb,
        peak_memory_mb: peak_mb,
        weight_memory_mb,
        activation_memory_mb,
        parameter_count,
        inference_latency_us,
        memory_per_parameter_bytes: (weight_memory_mb * 1_048_576.0) / parameter_count as f64,
    })
}

/// Estimate TFT parameter count
fn estimate_tft_parameters(config: &TFTConfig) -> usize {
    // Variable selection networks
    let vsn_params = 3 * (config.hidden_dim * config.hidden_dim);

    // Encoder/decoder stacks (GRN)
    let grn_params = 3 * config.num_layers * (config.hidden_dim * config.hidden_dim);

    // Attention
    let attention_params = config.num_heads * (config.hidden_dim * config.hidden_dim);

    // Quantile output
    let quantile_params = config.hidden_dim * config.prediction_horizon * config.num_quantiles;

    vsn_params + grn_params + attention_params + quantile_params
}

/// Estimate Liquid network parameter count
fn estimate_liquid_parameters(config: &LiquidNetworkConfig) -> usize {
    let layer_params = config.input_dim * config.hidden_dim +
                      config.hidden_dim * config.hidden_dim * (config.num_layers - 1) +
                      config.hidden_dim * config.output_dim;
    layer_params
}

/// Print formatted profile
fn print_profile(profile: &ModelMemoryProfile) {
    println!("\n{} Model Memory Profile:", profile.model_name);
    println!("  Base Memory:       {:>8.2} MB", profile.base_memory_mb);
    println!("  Weight Memory:     {:>8.2} MB", profile.weight_memory_mb);
    println!("  Activation Memory: {:>8.2} MB", profile.activation_memory_mb);
    println!("  Peak Memory:       {:>8.2} MB", profile.peak_memory_mb);
    println!("  Parameter Count:   {:>8}", profile.parameter_count);
    println!("  Bytes/Parameter:   {:>8.2}", profile.memory_per_parameter_bytes);
    println!("  Inference Latency: {:>8} µs", profile.inference_latency_us);

    // Check against targets
    let target_mb = match profile.model_name.as_str() {
        "DQN" => 256.0,
        "PPO" => 384.0,
        "TFT" => 512.0,
        "MAMBA-2" => 512.0,
        "Liquid" => 256.0,
        _ => 512.0,
    };

    let status = if profile.peak_memory_mb <= target_mb {
        "✅ MEETS TARGET"
    } else {
        "⚠️ EXCEEDS TARGET"
    };

    println!("  Target:            {:>8.2} MB", target_mb);
    println!("  Status:            {}", status);
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== ML Model Memory Profiling Tool ===\n");
    println!("Measuring memory usage for all models...\n");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Using device: {:?}\n", device);

    // Profile all models
    let mut profiles = Vec::new();

    println!("Profiling DQN...");
    match profile_dqn(&device) {
        Ok(profile) => {
            print_profile(&profile);
            profiles.push(profile);
        }
        Err(e) => eprintln!("Failed to profile DQN: {}", e),
    }

    println!("\nProfiling PPO...");
    match profile_ppo(&device) {
        Ok(profile) => {
            print_profile(&profile);
            profiles.push(profile);
        }
        Err(e) => eprintln!("Failed to profile PPO: {}", e),
    }

    println!("\nProfiling TFT...");
    match profile_tft(&device) {
        Ok(profile) => {
            print_profile(&profile);
            profiles.push(profile);
        }
        Err(e) => eprintln!("Failed to profile TFT: {}", e),
    }

    println!("\nProfiling MAMBA-2...");
    match profile_mamba(&device) {
        Ok(profile) => {
            print_profile(&profile);
            profiles.push(profile);
        }
        Err(e) => eprintln!("Failed to profile MAMBA-2: {}", e),
    }

    println!("\nProfiling Liquid...");
    match profile_liquid() {
        Ok(profile) => {
            print_profile(&profile);
            profiles.push(profile);
        }
        Err(e) => eprintln!("Failed to profile Liquid: {}", e),
    }

    // Summary
    println!("\n=== Summary ===\n");
    let total_memory: f64 = profiles.iter().map(|p| p.peak_memory_mb).sum();
    let total_params: usize = profiles.iter().map(|p| p.parameter_count).sum();

    println!("Total Memory (all models): {:.2} MB", total_memory);
    println!("Total Parameters:          {}", total_params);
    println!("Average Memory/Model:      {:.2} MB", total_memory / profiles.len() as f64);

    // Identify optimization opportunities
    println!("\n=== Optimization Opportunities ===\n");
    for profile in &profiles {
        let target_mb = match profile.model_name.as_str() {
            "DQN" => 256.0,
            "PPO" => 384.0,
            "TFT" => 512.0,
            "MAMBA-2" => 512.0,
            "Liquid" => 256.0,
            _ => 512.0,
        };

        if profile.peak_memory_mb > target_mb {
            let excess = profile.peak_memory_mb - target_mb;
            let reduction_needed = (excess / profile.peak_memory_mb) * 100.0;
            println!("{}: Needs {:.0}% reduction ({:.2} MB excess)",
                    profile.model_name, reduction_needed, excess);

            // Specific recommendations
            if profile.memory_per_parameter_bytes > 6.0 {
                println!("  → Convert to float16 (50% reduction expected)");
            }
            if profile.activation_memory_mb > profile.weight_memory_mb * 0.5 {
                println!("  → Implement gradient checkpointing");
            }
            if profile.parameter_count > 1_000_000 {
                println!("  → Apply 8-bit quantization");
            }
        }
    }

    // Save baseline report
    let report = MemoryOptimizationReport {
        timestamp: chrono::Utc::now().to_rfc3339(),
        baseline_profiles: profiles.clone(),
        optimized_profiles: Vec::new(), // To be filled after optimizations
        memory_savings_mb: HashMap::new(),
        memory_reduction_percent: HashMap::new(),
        accuracy_impact_percent: HashMap::new(),
        recommendations: vec![
            "Implement lazy checkpoint loading".to_string(),
            "Add float16 precision for inference".to_string(),
            "Implement 8-bit weight quantization".to_string(),
            "Use gradient checkpointing for training".to_string(),
            "Add model pruning for production deployment".to_string(),
        ],
    };

    let report_json = serde_json::to_string_pretty(&report)?;
    std::fs::write("memory_baseline_profile.json", report_json)?;
    println!("\nBaseline profile saved to memory_baseline_profile.json");

    Ok(())
}