#!/usr/bin/env rust-script
//! Comprehensive demo of LoRA for Diffusion Models
//! 
//! This demonstrates the complete LoRA implementation with all features:
//! - Different adapter configurations for various use cases
//! - Multiple adapter support with blending
//! - Memory-efficient parameter management
//! - Layer targeting and selective adaptation
//! - Advanced adapter management utilities

use rtx_diffuse::{
    LoRAConfig, LoRAAdapter, AttentionLoRALayer, LayerTargeting, 
    LoRAUNet, LoRAManager, UNetConfig, DiffusionError, Result
};
use rtx_tensor::{Device, Tensor};

fn main() -> Result<()> {
    println!("🎨 LoRA for Diffusion Models - Comprehensive Demo");
    println!("================================================\n");
    
    let device = Device::cuda(0).unwrap_or(Device::default());
    
    // 1. Different LoRA configurations for various use cases
    println!("1. LoRA Configuration Presets:");
    
    let style_config = LoRAConfig::for_style_transfer(16);
    println!("   Style Transfer: rank={}, alpha={}, scaling={:.2}", 
        style_config.rank, style_config.alpha, style_config.scaling());
    
    let concept_config = LoRAConfig::for_concept_learning(8);  
    println!("   Concept Learning: rank={}, alpha={}, scaling={:.2}",
        concept_config.rank, concept_config.alpha, concept_config.scaling());
    
    let full_config = LoRAConfig::for_full_fine_tuning(32);
    println!("   Full Fine-tuning: rank={}, alpha={}, scaling={:.2}",
        full_config.rank, full_config.alpha, full_config.scaling());
    
    // 2. Create adapters for different art styles
    println!("\n2. Creating Style Adapters:");
    
    let anime_adapter = LoRAAdapter::new("anime_style", style_config.clone(), &device)?;
    println!("   ✓ Anime Style: {} parameters", anime_adapter.total_parameters());
    
    let realistic_adapter = LoRAAdapter::new("realistic_portrait", concept_config, &device)?;
    println!("   ✓ Realistic Portrait: {} parameters", realistic_adapter.total_parameters());
    
    let abstract_adapter = LoRAAdapter::new("abstract_art", full_config, &device)?;
    println!("   ✓ Abstract Art: {} parameters", abstract_adapter.total_parameters());
    
    // 3. Attention layer targeting
    println!("\n3. Layer Targeting:");
    
    let targeting = LayerTargeting::new(vec![
        "self_attn".to_string(),
        "cross_attn".to_string(),
        "to_q".to_string(),
        "to_k".to_string(),
        "to_v".to_string()
    ]);
    
    let test_layers = [
        "encoder.0.self_attn.to_q",
        "encoder.0.cross_attn.to_k", 
        "decoder.mlp.fc1",
        "encoder.1.self_attn.to_v"
    ];
    
    for layer in &test_layers {
        let should_apply = targeting.should_apply_lora(layer);
        println!("   {} -> {}", layer, if should_apply { "✓ Apply" } else { "✗ Skip" });
    }
    
    // 4. UNet with multiple LoRA adapters
    println!("\n4. UNet LoRA Integration:");
    
    let unet_config = UNetConfig::default();
    let mut lora_unet = LoRAUNet::new(unet_config, &device)?;
    
    // Apply adapters
    lora_unet.apply_lora_adapter(anime_adapter)?;
    lora_unet.apply_lora_adapter(realistic_adapter)?;
    lora_unet.apply_lora_adapter(abstract_adapter)?;
    
    println!("   ✓ Applied {} adapters", lora_unet.num_lora_adapters());
    println!("   ✓ Total adapter parameters: {}", lora_unet.total_adapter_parameters());
    println!("   ✓ Total adapter memory: {} MB", 
        lora_unet.total_adapter_memory() / (1024 * 1024));
    
    // 5. Adapter blending and control
    println!("\n5. Adapter Blending:");
    
    // Blend adapters: 50% anime, 30% realistic, 20% abstract
    lora_unet.set_adapter_weights(vec![0.5, 0.3, 0.2])?;
    println!("   ✓ Set blend weights: [0.5, 0.3, 0.2]");
    
    // Dynamically enable/disable adapters
    lora_unet.set_adapter_enabled("abstract_art", false)?;
    println!("   ✓ Disabled abstract art adapter");
    
    println!("   Active adapters: {:?}", lora_unet.adapter_names());
    
    // 6. Advanced adapter management
    println!("\n6. Advanced LoRA Management:");
    
    let mut manager = LoRAManager::new(&device);
    
    // Create adapter library
    let styles = [
        ("cyberpunk", 16),
        ("watercolor", 8),
        ("oil_painting", 24),
        ("sketch", 4)
    ];
    
    manager.create_adapter_library(&styles)?;
    println!("   ✓ Created adapter library with {} styles", styles.len());
    
    // Optimal rank suggestion
    let optimal_rank = manager.suggest_optimal_rank(100.0, 512); // 100MB budget, 512 dim
    println!("   ✓ Suggested optimal rank for 100MB budget: {}", optimal_rank);
    
    // 7. Memory efficiency analysis
    println!("\n7. Memory Efficiency Analysis:");
    
    // Create sample attention layer for analysis
    let layer = AttentionLoRALayer::new(
        "sample_attention", 
        512, 
        512, 
        LoRAConfig::default(),
        &device
    )?;
    
    println!("   Layer: {} -> {}", layer.in_features(), layer.out_features());
    println!("   LoRA parameters: {} (rank {})", layer.num_parameters(), 8);
    println!("   Compression ratio: {:.1}x", layer.compression_ratio());
    println!("   Memory savings: {:.1}%", 
        (1.0 - (1.0 / layer.compression_ratio())) * 100.0);
    
    // 8. Parameter shapes and initialization
    println!("\n8. Parameter Analysis:");
    
    let (a_shape, b_shape) = layer.parameter_shapes();
    println!("   Matrix A shape: {:?}", a_shape);
    println!("   Matrix B shape: {:?}", b_shape);
    println!("   Dropout rate: {}", layer.dropout_rate());
    println!("   Spectral norm: {}", layer.has_spectral_norm());
    println!("   Gradient checkpointing: {}", layer.has_gradient_checkpointing());
    
    // 9. Forward pass demonstration
    println!("\n9. Forward Pass Test:");
    
    // Create test tensors
    let input = create_test_tensor(&[2, 64, 512], &device)?;  // [batch, seq, dim]
    let weight = create_test_tensor(&[512, 512], &device)?;   // [out, in]
    
    let output = layer.forward(&input, &weight)?;
    println!("   ✓ Forward pass: {:?} -> {:?}", 
        input.shape().dims(), output.shape().dims());
    
    // 10. Weight merging demonstration
    println!("\n10. Weight Merging:");
    
    let merged = layer.merge_weights(&weight)?;
    println!("   ✓ Merged weights shape: {:?}", merged.shape().dims());
    
    let unmerged = layer.unmerge_weights(&merged)?;
    println!("   ✓ Unmerged weights shape: {:?}", unmerged.shape().dims());
    
    println!("\n🎉 LoRA for Diffusion Demo Complete!");
    println!("\n📊 Summary of Features Demonstrated:");
    println!("   • LoRA adapter creation and management");
    println!("   • Multiple adapter support with weighted blending");
    println!("   • Selective layer targeting for efficient fine-tuning");
    println!("   • Memory-efficient low-rank adaptation (up to 100x compression)");
    println!("   • Different configurations for style transfer, concept learning, etc.");
    println!("   • Advanced adapter fusion and management utilities");
    println!("   • Proper initialization and forward pass implementation");
    println!("   • Weight merging for inference optimization");
    
    Ok(())
}

// Helper function to create test tensors
fn create_test_tensor(shape: &[usize], device: &Device) -> Result<Tensor> {
    let size = shape.iter().product::<usize>();
    let data = vec![1.0f32; size];
    
    Tensor::new(data, shape.to_vec()).map_err(|e| DiffusionError::TensorOperation {
        details: format!("Failed to create test tensor: {:?}", e),
    })
}