//! Quantized Checkpoint Demo
//!
//! Demonstrates saving and loading quantized INT8 model checkpoints with SafeTensors.
//!
//! Usage:
//! ```bash
//! cargo run --example quantized_checkpoint_demo --release
//! ```
//!
//! Features:
//! - Create synthetic DQN model weights
//! - Quantize FP32 → INT8 (4x compression)
//! - Save to SafeTensors format
//! - Load and validate round-trip
//! - Compare file sizes (FP32 vs INT8)

use candle_core::{DType, Device, Tensor};
use ml::checkpoint::{
    calculate_compression_ratio, load_quantized_checkpoint, save_quantized_checkpoint,
    QuantizedCheckpointMetadata, QuantizedWeight,
};
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::MLError;
use std::collections::HashMap;
use std::path::PathBuf;
use tracing::{info, Level};
use tracing_subscriber;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    tracing_subscriber::fmt().with_max_level(Level::INFO).init();

    info!("=== Quantized Checkpoint Demo ===");

    // Create device
    let device = Device::Cpu;

    // Step 1: Create synthetic FP32 model weights (simulating DQN)
    info!("\n[Step 1] Creating synthetic DQN model weights (FP32)...");

    let mut fp32_weights: HashMap<String, Tensor> = HashMap::new();

    // Input layer: 225 features × 128 hidden
    let fc1_weight = create_random_tensor(&[225, 128], &device)?;
    fp32_weights.insert("fc1.weight".to_string(), fc1_weight);

    let fc1_bias = create_random_tensor(&[128], &device)?;
    fp32_weights.insert("fc1.bias".to_string(), fc1_bias);

    // Hidden layer: 128 × 64
    let fc2_weight = create_random_tensor(&[128, 64], &device)?;
    fp32_weights.insert("fc2.weight".to_string(), fc2_weight);

    let fc2_bias = create_random_tensor(&[64], &device)?;
    fp32_weights.insert("fc2.bias".to_string(), fc2_bias);

    // Output layer: 64 × 3 (BUY/SELL/HOLD)
    let fc3_weight = create_random_tensor(&[64, 3], &device)?;
    fp32_weights.insert("fc3.weight".to_string(), fc3_weight);

    let fc3_bias = create_random_tensor(&[3], &device)?;
    fp32_weights.insert("fc3.bias".to_string(), fc3_bias);

    let total_params: usize = fp32_weights
        .values()
        .map(|t| t.dims().iter().product::<usize>())
        .sum();
    let fp32_size = total_params * 4; // 4 bytes per F32

    info!(
        "Created {} layers, {} parameters ({:.2} MB FP32)",
        fp32_weights.len(),
        total_params,
        fp32_size as f64 / 1_048_576.0
    );

    // Step 2: Quantize weights to INT8
    info!("\n[Step 2] Quantizing weights to INT8...");

    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };

    let mut quantizer = Quantizer::new(quant_config, device.clone());
    let mut quantized_weights: HashMap<String, QuantizedWeight> = HashMap::new();

    for (name, tensor) in &fp32_weights {
        let quantized_tensor = quantizer.quantize_tensor(tensor, name)?;
        let quantized_weight = QuantizedWeight::from_quantized_tensor(&quantized_tensor)?;
        quantized_weights.insert(name.clone(), quantized_weight);
    }

    let int8_size: usize = quantized_weights.values().map(|w| w.memory_bytes()).sum();
    let compression_ratio = calculate_compression_ratio(&quantized_weights);

    info!(
        "Quantized to INT8: {:.2} MB ({:.2}x compression)",
        int8_size as f64 / 1_048_576.0,
        compression_ratio
    );

    // Step 3: Save quantized checkpoint
    info!("\n[Step 3] Saving quantized checkpoint...");

    let checkpoint_path = PathBuf::from("ml/trained_models/dqn_quantized_demo.safetensors");

    let metadata = QuantizedCheckpointMetadata {
        model_type: "DQN".to_string(),
        version: "1.0.0".to_string(),
        quantization_method: "symmetric".to_string(),
        quantization_type: "int8".to_string(),
        ..Default::default()
    };

    let file_size_uncompressed = save_quantized_checkpoint(
        &checkpoint_path,
        &quantized_weights,
        Some(metadata.clone()),
        false,
    )?;

    info!(
        "Saved checkpoint: {} ({:.2} MB uncompressed)",
        checkpoint_path.display(),
        file_size_uncompressed as f64 / 1_048_576.0
    );

    // Step 4: Save compressed version
    let compressed_path =
        PathBuf::from("ml/trained_models/dqn_quantized_demo_compressed.safetensors");
    let file_size_compressed =
        save_quantized_checkpoint(&compressed_path, &quantized_weights, Some(metadata), true)?;

    info!(
        "Saved compressed: {} ({:.2} MB, {:.1}% reduction)",
        compressed_path.display(),
        file_size_compressed as f64 / 1_048_576.0,
        100.0 * (1.0 - file_size_compressed as f64 / file_size_uncompressed as f64)
    );

    // Step 5: Load checkpoint and validate
    info!("\n[Step 4] Loading checkpoint and validating...");

    let (loaded_weights, loaded_metadata) = load_quantized_checkpoint(&checkpoint_path)?;

    info!("Loaded {} layers from checkpoint", loaded_weights.len());
    info!(
        "Metadata: model_type={}, version={}, num_layers={}",
        loaded_metadata.model_type, loaded_metadata.version, loaded_metadata.num_layers
    );

    // Validate weights match
    let mut all_match = true;
    for (name, original) in &quantized_weights {
        if let Some(loaded) = loaded_weights.get(name) {
            if original.scale != loaded.scale || original.zero_point != loaded.zero_point {
                info!("❌ Mismatch in {}: scale or zero_point differs", name);
                all_match = false;
            }
        } else {
            info!("❌ Missing layer: {}", name);
            all_match = false;
        }
    }

    if all_match {
        info!("✅ All weights validated successfully!");
    }

    // Step 6: Dequantize and compare
    info!("\n[Step 5] Dequantizing and comparing with original FP32...");

    let mut max_error = 0.0f32;
    let mut avg_error = 0.0f32;
    let mut error_count = 0;

    for (name, original_fp32) in &fp32_weights {
        if let Some(loaded_weight) = loaded_weights.get(name) {
            // Dequantize: x = scale * (q - zero_point)
            let quantized_tensor = loaded_weight.to_quantized_tensor();
            let dequantized = quantizer.dequantize_tensor(&quantized_tensor)?;

            // Compare with original
            let original_vec = original_fp32.flatten_all()?.to_vec1::<f32>()?;
            let dequant_vec = dequantized.flatten_all()?.to_vec1::<f32>()?;

            for (orig, dequant) in original_vec.iter().zip(dequant_vec.iter()) {
                let error = (orig - dequant).abs();
                max_error = max_error.max(error);
                avg_error += error;
                error_count += 1;
            }
        }
    }

    avg_error /= error_count as f32;

    info!(
        "Quantization error: max={:.6}, avg={:.6}",
        max_error, avg_error
    );

    // Step 7: Size comparison summary
    info!("\n=== Size Comparison Summary ===");
    info!(
        "FP32 (original):     {:.2} MB",
        fp32_size as f64 / 1_048_576.0
    );
    info!(
        "INT8 (quantized):    {:.2} MB",
        int8_size as f64 / 1_048_576.0
    );
    info!(
        "File (uncompressed): {:.2} MB",
        file_size_uncompressed as f64 / 1_048_576.0
    );
    info!(
        "File (compressed):   {:.2} MB",
        file_size_compressed as f64 / 1_048_576.0
    );
    info!("");
    info!("Compression ratio: {:.2}x (FP32 → INT8)", compression_ratio);
    info!(
        "Total savings: {:.2} MB ({:.1}%)",
        (fp32_size - file_size_compressed) as f64 / 1_048_576.0,
        100.0 * (1.0 - file_size_compressed as f64 / fp32_size as f64)
    );

    info!("\n✅ Demo complete! Checkpoints saved to:");
    info!("   - {}", checkpoint_path.display());
    info!("   - {}", compressed_path.display());

    Ok(())
}

/// Create random FP32 tensor with values in [-1.0, 1.0]
fn create_random_tensor(shape: &[usize], device: &Device) -> Result<Tensor, MLError> {
    use rand::Rng;
    let mut rng = rand::thread_rng();

    let total_elements: usize = shape.iter().product();
    let data: Vec<f32> = (0..total_elements)
        .map(|_| rng.gen_range(-1.0..1.0))
        .collect();

    Tensor::from_vec(data, shape, device)
        .map_err(|e| MLError::ModelError(format!("Failed to create tensor: {}", e)))
}