//! Comprehensive model compression example
//!
//! This example demonstrates how to use the rtx-compress crate to compress
//! a neural network model using various techniques:
//! - Post-training quantization
//! - Structured and unstructured pruning
//! - Knowledge distillation
//! - Integrated compression pipeline

use rtx_compress::{
    Result,
    // Distillation
    distillation::{
        DistillationConfig, DistillationLoss, DistillationMethod, FeatureMatchingConfig,
        KnowledgeDistiller,
    },
    // Pipeline
    pipeline::{
        CompressionPipeline, CompressionPipelineConfig, CompressionStrategy, HardwareTarget,
    },
    // Pruning
    pruning::{
        GradualPruningSchedule, ImportanceMetric, PruningCriterion,
        PruningMethod as StructuredPruningMethod, StructuredPruner, StructuredPruningConfig,
        UnstructuredPruner, UnstructuredPruningConfig,
    },
    // Quantization
    quantization::{
        CalibrationConfig, CalibrationMethod, PostTrainingQuantizer, QuantizationConfig,
        QuantizationScheme,
    },
};
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

/// Create a simple test model with realistic parameter distributions
fn create_test_model() -> Result<HashMap<String, Tensor>> {
    let device = Device::try_default()?;
    let mut model = HashMap::new();

    // Convolutional layers
    model.insert(
        "conv1.weight".to_string(),
        create_conv_weights(&device, 32, 3, 5, 5)?,
    );
    model.insert("conv1.bias".to_string(), create_bias(&device, 32)?);

    model.insert(
        "conv2.weight".to_string(),
        create_conv_weights(&device, 64, 32, 5, 5)?,
    );
    model.insert("conv2.bias".to_string(), create_bias(&device, 64)?);

    model.insert(
        "conv3.weight".to_string(),
        create_conv_weights(&device, 128, 64, 3, 3)?,
    );
    model.insert("conv3.bias".to_string(), create_bias(&device, 128)?);

    // Fully connected layers
    model.insert(
        "fc1.weight".to_string(),
        create_fc_weights(&device, 512, 2048)?,
    );
    model.insert("fc1.bias".to_string(), create_bias(&device, 512)?);

    model.insert(
        "fc2.weight".to_string(),
        create_fc_weights(&device, 256, 512)?,
    );
    model.insert("fc2.bias".to_string(), create_bias(&device, 256)?);

    model.insert(
        "fc3.weight".to_string(),
        create_fc_weights(&device, 10, 256)?,
    );
    model.insert("fc3.bias".to_string(), create_bias(&device, 10)?);

    Ok(model)
}

fn create_conv_weights(
    device: &Device,
    out_channels: usize,
    in_channels: usize,
    h: usize,
    w: usize,
) -> Result<Tensor> {
    let total = out_channels * in_channels * h * w;
    let mut weights = Vec::with_capacity(total);

    for i in 0..total {
        // Create realistic weight distribution (Xavier initialization)
        let fan_in = in_channels * h * w;
        let scale = (2.0 / fan_in as f32).sqrt();
        let val = scale * ((i % 1000) as f32 / 500.0 - 1.0);
        weights.push(val);
    }

    Tensor::from_slice(&weights, &[out_channels, in_channels, h, w], device)
}

fn create_fc_weights(device: &Device, out_features: usize, in_features: usize) -> Result<Tensor> {
    let total = out_features * in_features;
    let mut weights = Vec::with_capacity(total);

    for i in 0..total {
        // Xavier initialization for fully connected
        let scale = (2.0 / in_features as f32).sqrt();
        let val = scale * ((i % 1000) as f32 / 500.0 - 1.0);
        weights.push(val);
    }

    Tensor::from_slice(&weights, &[out_features, in_features], device)
}

fn create_bias(device: &Device, size: usize) -> Result<Tensor> {
    let bias = vec![0.01; size]; // Small positive bias
    Tensor::from_slice(&bias, &[size], device)
}

fn calculate_model_size(model: &HashMap<String, Tensor>) -> (usize, f32) {
    let total_params: usize = model.values().map(|t| t.numel()).sum();
    let size_mb = (total_params * 4) as f32 / (1024.0 * 1024.0);
    (total_params, size_mb)
}

fn print_model_info(name: &str, model: &HashMap<String, Tensor>) {
    let (params, size_mb) = calculate_model_size(model);
    println!("{}: {} parameters, {:.2} MB", name, params, size_mb);
}

fn demonstrate_post_training_quantization() -> Result<()> {
    println!("\n=== Post-Training Quantization Demo ===");

    let model = create_test_model()?;
    print_model_info("Original Model", &model);

    // Create calibration dataset
    let device = Device::try_default()?;
    let mut calibration_data = Vec::new();
    for _ in 0..100 {
        calibration_data.push(Tensor::randn(&[32, 64], &device)?);
    }

    // INT8 quantization
    let int8_config =
        QuantizationConfig::new(QuantizationScheme::INT8, CalibrationConfig::new(100, 0.99));

    let mut int8_quantizer = PostTrainingQuantizer::new(int8_config)?;
    int8_quantizer.calibrate(&calibration_data)?;

    println!("INT8 Quantization Statistics:");
    let stats = int8_quantizer.get_statistics();
    println!("  Scale: {:.6}", stats.scale);
    println!("  Zero Point: {}", stats.zero_point);
    println!("  Bit Width: {}", stats.bit_width);

    // INT4 quantization for more aggressive compression
    let int4_config =
        QuantizationConfig::new(QuantizationScheme::INT4, CalibrationConfig::new(100, 0.95));

    let mut int4_quantizer = PostTrainingQuantizer::new(int4_config)?;
    int4_quantizer.calibrate(&calibration_data)?;

    let int4_stats = int4_quantizer.get_statistics();
    println!("INT4 Quantization Statistics:");
    println!("  Scale: {:.6}", int4_stats.scale);
    println!("  Zero Point: {}", int4_stats.zero_point);
    println!("  Expected compression: ~8x");

    Ok(())
}

fn demonstrate_structured_pruning() -> Result<()> {
    println!("\n=== Structured Pruning Demo ===");

    let model = create_test_model()?;
    print_model_info("Original Model", &model);

    // Channel pruning with L2 importance
    let config = StructuredPruningConfig::new(
        StructuredPruningMethod::ChannelPruning,
        ImportanceMetric::L2Norm,
        0.5, // 50% sparsity
    );

    let pruner = StructuredPruner::new(config)?;

    let mut pruned_model = HashMap::new();
    for (name, tensor) in &model {
        let pruned_tensor = pruner.prune_tensor(tensor, name)?;
        pruned_model.insert(name.clone(), pruned_tensor);
    }

    print_model_info("Structured Pruned Model", &pruned_model);

    let (orig_params, _) = calculate_model_size(&model);
    let (pruned_params, _) = calculate_model_size(&pruned_model);
    let compression_ratio = orig_params as f32 / pruned_params as f32;

    println!("Structured Pruning Results:");
    println!("  Compression Ratio: {:.2}x", compression_ratio);
    println!("  Parameters Removed: {}", orig_params - pruned_params);

    Ok(())
}

fn demonstrate_unstructured_pruning() -> Result<()> {
    println!("\n=== Unstructured Pruning Demo ===");

    let model = create_test_model()?;

    // Global magnitude pruning
    let global_config = UnstructuredPruningConfig::new(
        PruningCriterion::GlobalMagnitude,
        0.8, // 80% sparsity
    );

    let global_pruner = UnstructuredPruner::new(global_config)?;
    let globally_pruned = global_pruner.prune_model(&model)?;

    print_model_info("Globally Pruned Model (80% sparse)", &globally_pruned);

    // Gradual pruning demonstration
    let schedule = GradualPruningSchedule::polynomial(
        0.0, // Start with no pruning
        0.9, // End with 90% sparsity
        100, // Over 100 steps
        3.0, // Cubic schedule
    );

    let gradual_config =
        UnstructuredPruningConfig::new_with_schedule(PruningCriterion::GlobalMagnitude, schedule);

    let mut gradual_pruner = UnstructuredPruner::new(gradual_config)?;

    // Simulate gradual pruning over training
    let mut current_model = model.clone();
    let steps = [0, 25, 50, 75, 99];

    println!("Gradual Pruning Progress:");
    for &step in &steps {
        gradual_pruner.set_current_step(step);
        current_model = gradual_pruner.prune_model(&current_model)?;

        let (params, _) = calculate_model_size(&current_model);
        let (orig_params, _) = calculate_model_size(&model);
        let sparsity = 1.0 - (params as f32 / orig_params as f32);

        println!("  Step {}: {:.1}% sparse", step, sparsity * 100.0);
    }

    Ok(())
}

fn demonstrate_knowledge_distillation() -> Result<()> {
    println!("\n=== Knowledge Distillation Demo ===");

    let device = Device::try_default()?;

    // Simulate teacher and student logits
    let batch_size = 32;
    let num_classes = 10;

    let teacher_logits = Tensor::randn(&[batch_size, num_classes], &device)?;
    let student_logits = Tensor::randn(&[batch_size, num_classes], &device)?;
    let targets = Tensor::from_slice(
        &(0..batch_size)
            .map(|i| (i % num_classes) as i64)
            .collect::<Vec<_>>(),
        &[batch_size],
        &device,
    )?;

    // Response-based distillation
    let response_config = DistillationConfig::new(
        DistillationMethod::ResponseBased {
            temperature: 4.0,
            alpha: 0.7,
        },
        DistillationLoss::KullbackLeibler,
    );

    let distiller = KnowledgeDistiller::new(response_config)?;

    let result = distiller.compute_distillation_loss(
        &student_logits,
        &teacher_logits,
        Some(&targets),
        None,
    )?;

    println!("Response-based Distillation:");
    println!("  Total Loss: {:.4}", result.total_loss);
    println!("  Distillation Loss: {:.4}", result.distillation_loss);
    if let Some(task_loss) = result.task_loss {
        println!("  Task Loss: {:.4}", task_loss);
    }

    // Feature-based distillation
    let feature_config = FeatureMatchingConfig::new(
        vec![128, 256, 512], // Feature dimensions
        vec![0.2, 0.3, 0.5], // Layer weights
        "mse".to_string(),
    );

    let feature_method = DistillationMethod::FeatureBased {
        feature_config,
        intermediate_matching: true,
    };

    let feature_distiller = KnowledgeDistiller::new(DistillationConfig::new(
        feature_method,
        DistillationLoss::MeanSquaredError,
    ))?;

    // Simulate intermediate features
    let teacher_features = vec![
        Tensor::randn(&[batch_size, 128], &device)?,
        Tensor::randn(&[batch_size, 256], &device)?,
        Tensor::randn(&[batch_size, 512], &device)?,
    ];

    let student_features = vec![
        Tensor::randn(&[batch_size, 128], &device)?,
        Tensor::randn(&[batch_size, 256], &device)?,
        Tensor::randn(&[batch_size, 512], &device)?,
    ];

    let feature_result = feature_distiller
        .compute_feature_distillation_loss(&student_features, &teacher_features)?;

    println!("Feature-based Distillation:");
    println!("  Total Loss: {:.4}", feature_result.total_loss);
    println!("  Layer Losses: {:?}", feature_result.feature_losses);

    Ok(())
}

fn demonstrate_compression_pipeline() -> Result<()> {
    println!("\n=== Comprehensive Compression Pipeline Demo ===");

    let model = create_test_model()?;
    print_model_info("Original Model", &model);

    // Automatic pipeline for mobile deployment
    let mut mobile_pipeline = CompressionPipeline::auto(
        4.0,  // 4x compression target
        0.95, // 95% accuracy retention
        HardwareTarget::Mobile,
    )?;

    let compressed_result = mobile_pipeline.compress(&model, None, None)?;

    print_model_info("Compressed Model", &compressed_result.compressed_model);

    println!("Compression Results:");
    println!(
        "  Compression Ratio: {:.2}x",
        compressed_result.statistics.compression_ratio
    );
    println!(
        "  Size Reduction: {:.2} MB",
        compressed_result.statistics.size_reduction_mb
    );
    println!(
        "  Estimated Speedup: {:.2}x",
        compressed_result.statistics.estimated_speedup
    );

    // Custom strategy for high compression
    let custom_strategy = CompressionStrategy::Custom {
        techniques: vec![
            rtx_compress::pipeline::CompressionTechnique::UnstructuredPruning,
            rtx_compress::pipeline::CompressionTechnique::StructuredPruning,
            rtx_compress::pipeline::CompressionTechnique::Quantization,
        ],
        priorities: [
            (
                rtx_compress::pipeline::CompressionTechnique::UnstructuredPruning,
                0.4,
            ),
            (
                rtx_compress::pipeline::CompressionTechnique::StructuredPruning,
                0.3,
            ),
            (
                rtx_compress::pipeline::CompressionTechnique::Quantization,
                0.3,
            ),
        ]
        .iter()
        .cloned()
        .collect(),
    };

    let custom_config = CompressionPipelineConfig {
        strategy: custom_strategy,
        target_compression_ratio: 8.0,
        target_accuracy_retention: 0.90,
        progressive: true,
        num_stages: 3,
        validation_size: 1000,
        hardware_target: HardwareTarget::Generic,
    };

    let mut custom_pipeline = CompressionPipeline::new(custom_config)?;
    let custom_result = custom_pipeline.compress(&model, None, None)?;

    print_model_info("Custom Compressed Model", &custom_result.compressed_model);

    println!("Custom Compression Results:");
    println!(
        "  Compression Ratio: {:.2}x",
        custom_result.statistics.compression_ratio
    );
    println!(
        "  Size Reduction: {:.2} MB",
        custom_result.statistics.size_reduction_mb
    );

    Ok(())
}

fn main() -> Result<()> {
    println!("RustyTorch++ Comprehensive Model Compression Demo");
    println!("================================================");

    demonstrate_post_training_quantization()?;
    demonstrate_structured_pruning()?;
    demonstrate_unstructured_pruning()?;
    demonstrate_knowledge_distillation()?;
    demonstrate_compression_pipeline()?;

    println!("\n=== Summary ===");
    println!("This demo showcased comprehensive model compression techniques:");
    println!("1. Post-training quantization (INT8/INT4) for memory and speed");
    println!("2. Structured pruning for hardware-friendly compression");
    println!("3. Unstructured pruning for maximum compression ratios");
    println!("4. Knowledge distillation for accuracy preservation");
    println!("5. Integrated pipeline for automatic compression strategy");
    println!("\nAll techniques can be combined and customized for specific deployment needs.");

    Ok(())
}