// Standalone test for LoRA diffusion functionality

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

fn main() -> Result<()> {
    println!("Testing LoRA for Diffusion Models");
    
    let device = Device::cuda(0).unwrap_or(Device::default());
    
    // Test 1: LoRA Config and scaling
    let config = LoRAConfig {
        rank: 16,
        alpha: 32.0,
        ..Default::default()
    };
    println!("✓ LoRA config created with scaling: {}", config.alpha / config.rank as f32);
    
    // Test 2: Create LoRA adapter
    let adapter = LoRAAdapter::new("test_adapter", config.clone(), &device)?;
    println!("✓ LoRA adapter '{}' created with rank {}", adapter.name(), adapter.rank());
    
    // Test 3: Layer targeting 
    let targeting = LayerTargeting::new(vec!["attn".to_string(), "cross_attn".to_string()]);
    assert!(targeting.should_apply_lora("encoder.0.attn.q_proj"));
    assert!(!targeting.should_apply_lora("decoder.mlp.fc1"));
    println!("✓ Layer targeting works correctly");
    
    // Test 4: Attention LoRA layer
    let layer = AttentionLoRALayer::new(
        "test_layer",
        512,
        256,
        config,
        &device
    )?;
    
    println!("✓ Attention LoRA layer created:");
    println!("  - Input features: {}", layer.in_features());
    println!("  - Output features: {}", layer.out_features());
    println!("  - Parameters: {}", layer.num_parameters());
    println!("  - Compression ratio: {:.2}x", layer.compression_ratio());
    
    // Test 5: LoRA UNet integration
    let unet_config = UNetConfig::default();
    let mut lora_unet = LoRAUNet::new(unet_config, &device)?;
    
    lora_unet.apply_lora_adapter(adapter)?;
    assert!(lora_unet.has_lora_adapter("test_adapter"));
    assert_eq!(lora_unet.num_lora_adapters(), 1);
    println!("✓ LoRA UNet integration successful");
    
    // Test 6: Multiple adapters with weighting
    let adapter2 = LoRAAdapter::new("style2", LoRAConfig::default(), &device)?;
    lora_unet.apply_lora_adapter(adapter2)?;
    
    lora_unet.set_adapter_weights(vec![0.7, 0.3])?;
    assert_eq!(lora_unet.num_lora_adapters(), 2);
    println!("✓ Multiple adapters with weighting successful");
    
    println!("\n🎉 All LoRA diffusion tests passed!");
    println!("LoRA for Diffusion implementation complete with:");
    println!("- Low-rank adaptation for UNet attention layers");
    println!("- Configurable ranks (4, 8, 16, 32)"); 
    println!("- Alpha scaling factor control");
    println!("- Selective layer targeting");
    println!("- Weight merging/unmerging");
    println!("- Multiple adapter support");
    
    Ok(())
}