//! Tests for LoRA (Low-Rank Adaptation) functionality
//!
//! TDD: Define expected behavior for LoRA compression

use anyhow::Result;
use approx::assert_abs_diff_eq;
use rtx_compress::lora::{LoRACompressor, LoRAConfig, LoRALayer};
use rtx_tensor::{Device, Tensor};

#[test]
fn test_lora_config_creation() {
    let config = LoRAConfig {
        rank: 16,
        alpha: 32.0,
        dropout: 0.1,
        target_modules: vec![
            "q_proj".to_string(),
            "v_proj".to_string(),
            "k_proj".to_string(),
            "o_proj".to_string(),
        ],
        merge_weights: false,
    };

    assert_eq!(config.rank, 16);
    assert_eq!(config.alpha, 32.0);
    assert_eq!(config.scaling(), 32.0 / 16.0); // alpha / rank
}

#[test]
fn test_lora_layer_creation() {
    let device = Device::cpu();

    // Original weight dimensions
    let in_features = 768;
    let out_features = 768;

    let config = LoRAConfig {
        rank: 8,
        alpha: 16.0,
        dropout: 0.0,
        target_modules: vec![],
        merge_weights: false,
    };

    let lora_layer = LoRALayer::new(in_features, out_features, config, &device);
    assert!(lora_layer.is_ok());

    let lora = lora_layer.unwrap();

    // Check dimensions of LoRA matrices
    assert_eq!(lora.lora_a_shape(), vec![8, 768]); // [rank, in_features]
    assert_eq!(lora.lora_b_shape(), vec![768, 8]); // [out_features, rank]

    // Parameter reduction
    let original_params = in_features * out_features; // 768 * 768 = 589,824
    let lora_params = lora.num_parameters(); // 8*768 + 768*8 = 12,288
    assert_eq!(lora_params, 8 * 768 * 2);
    assert!(lora_params < original_params / 10); // Should be much smaller
}

#[test]
fn test_lora_forward_pass() {
    let device = Device::cpu();

    // Create original weight matrix
    let weight = Tensor::randn(&[256, 512], &device).unwrap();

    // Create LoRA decomposition
    let config = LoRAConfig {
        rank: 4,
        alpha: 8.0,
        dropout: 0.0,
        target_modules: vec![],
        merge_weights: false,
    };

    let lora = LoRALayer::new(512, 256, config, &device).unwrap();

    // Input
    let input = Tensor::randn(&[2, 512], &device).unwrap();

    // Forward pass through original weight
    let output_original = input.matmul(&weight.transpose(0, 1).unwrap()).unwrap();

    // Forward pass with LoRA
    let output_lora = lora.forward(&input, &weight).unwrap();

    // Shapes should match
    assert_eq!(output_lora.shape().dims(), output_original.shape().dims());
    assert_eq!(output_lora.shape().dims(), &[2, 256]);
}

#[test]
fn test_lora_weight_merging() {
    let device = Device::cpu();

    let weight = Tensor::randn(&[128, 128], &device).unwrap();

    let config = LoRAConfig {
        rank: 8,
        alpha: 16.0,
        dropout: 0.0,
        target_modules: vec![],
        merge_weights: true, // Enable merging
    };

    let mut lora = LoRALayer::new(128, 128, config, &device).unwrap();

    // Initialize LoRA matrices with specific values
    lora.initialize_lora_weights().unwrap();

    // Merge LoRA weights into original weight
    let merged = lora.merge_weights(&weight).unwrap();

    // Check dimensions
    assert_eq!(merged.shape().dims(), weight.shape().dims());

    // Unmerge should recover original weight
    let unmerged = lora.unmerge_weights(&merged).unwrap();

    let weight_data = weight.to_vec().unwrap();
    let unmerged_data = unmerged.to_vec().unwrap();

    for i in 0..weight_data.len() {
        assert_abs_diff_eq!(weight_data[i], unmerged_data[i], epsilon = 1e-4);
    }
}

#[test]
fn test_lora_compression_ratio() {
    let device = Device::cpu();

    // Test compression for different matrix sizes and ranks
    let test_cases = vec![
        (768, 768, 8),   // Transformer attention
        (768, 3072, 16), // Transformer FFN
        (1024, 1024, 4), // Large model, very low rank
    ];

    for (in_dim, out_dim, rank) in test_cases {
        let config = LoRAConfig {
            rank,
            alpha: rank as f32 * 2.0,
            dropout: 0.0,
            target_modules: vec![],
            merge_weights: false,
        };

        let lora = LoRALayer::new(in_dim, out_dim, config, &device).unwrap();

        let original_params = in_dim * out_dim;
        let lora_params = lora.num_parameters();
        let compression_ratio = lora.compression_ratio();

        assert_eq!(
            compression_ratio,
            original_params as f32 / lora_params as f32
        );
        assert!(compression_ratio > 1.0);

        // Verify parameter count
        assert_eq!(lora_params, rank * (in_dim + out_dim));
    }
}

#[test]
fn test_lora_adapter_training() {
    let device = Device::cpu();

    // Simulate fine-tuning scenario
    let base_weight = Tensor::randn(&[256, 256], &device).unwrap();

    let config = LoRAConfig {
        rank: 8,
        alpha: 16.0,
        dropout: 0.0,
        target_modules: vec![],
        merge_weights: false,
    };

    let mut lora = LoRALayer::new(256, 256, config, &device).unwrap();

    // Freeze base weight, only train LoRA parameters
    let trainable_params = lora.get_trainable_parameters();
    assert_eq!(trainable_params.len(), 2); // lora_A and lora_B

    // Simulate gradient update
    let grad_a = Tensor::randn(&[8, 256], &device).unwrap();
    let grad_b = Tensor::randn(&[256, 8], &device).unwrap();

    lora.update_parameters(&grad_a, &grad_b, 0.001).unwrap();

    // Parameters should have changed
    let new_params = lora.get_trainable_parameters();
    assert_eq!(new_params.len(), 2);
}

#[test]
fn test_multi_layer_lora_compression() {
    let device = Device::cpu();

    // Compress multiple layers with LoRA
    let compressor = LoRACompressor::new(
        LoRAConfig {
            rank: 16,
            alpha: 32.0,
            dropout: 0.0,
            target_modules: vec![
                "attention.q".to_string(),
                "attention.k".to_string(),
                "attention.v".to_string(),
            ],
            merge_weights: false,
        },
        &device,
    )
    .unwrap();

    // Original model weights
    let weights = vec![
        (
            "attention.q".to_string(),
            Tensor::randn(&[768, 768], &device).unwrap(),
        ),
        (
            "attention.k".to_string(),
            Tensor::randn(&[768, 768], &device).unwrap(),
        ),
        (
            "attention.v".to_string(),
            Tensor::randn(&[768, 768], &device).unwrap(),
        ),
        (
            "attention.o".to_string(),
            Tensor::randn(&[768, 768], &device).unwrap(),
        ),
        (
            "ffn.w1".to_string(),
            Tensor::randn(&[3072, 768], &device).unwrap(),
        ),
        (
            "ffn.w2".to_string(),
            Tensor::randn(&[768, 3072], &device).unwrap(),
        ),
    ];

    // Apply LoRA compression
    let compressed = compressor.compress_model(&weights).unwrap();

    // Check that only target modules are compressed
    assert_eq!(compressed.num_lora_layers(), 3);
    assert!(compressed.has_lora("attention.q"));
    assert!(compressed.has_lora("attention.k"));
    assert!(compressed.has_lora("attention.v"));
    assert!(!compressed.has_lora("attention.o"));

    // Calculate total compression
    let original_size = weights
        .iter()
        .filter(|(name, _)| compressor.is_target_module(name))
        .map(|(_, w)| w.numel())
        .sum::<usize>();

    let compressed_size = compressed.total_lora_parameters();
    let compression_ratio = original_size as f32 / compressed_size as f32;

    assert!(compression_ratio > 10.0); // Should achieve significant compression
}

#[test]
fn test_lora_with_different_ranks() {
    let device = Device::cpu();

    // Test that higher rank preserves more information
    let weight = Tensor::randn(&[512, 512], &device).unwrap();
    let input = Tensor::randn(&[4, 512], &device).unwrap();

    // Original output
    let output_original = input.matmul(&weight.transpose(0, 1).unwrap()).unwrap();

    // Low rank LoRA
    let lora_low = LoRALayer::new(
        512,
        512,
        LoRAConfig {
            rank: 4,
            alpha: 8.0,
            dropout: 0.0,
            target_modules: vec![],
            merge_weights: false,
        },
        &device,
    )
    .unwrap();

    // High rank LoRA
    let lora_high = LoRALayer::new(
        512,
        512,
        LoRAConfig {
            rank: 32,
            alpha: 64.0,
            dropout: 0.0,
            target_modules: vec![],
            merge_weights: false,
        },
        &device,
    )
    .unwrap();

    let output_low = lora_low.forward(&input, &weight).unwrap();
    let output_high = lora_high.forward(&input, &weight).unwrap();

    // Both should have correct shape
    assert_eq!(output_low.shape().dims(), output_original.shape().dims());
    assert_eq!(output_high.shape().dims(), output_original.shape().dims());

    // Higher rank should have more parameters
    assert!(lora_high.num_parameters() > lora_low.num_parameters());
}

#[test]
fn test_lora_initialization_methods() {
    let device = Device::cpu();

    let config = LoRAConfig {
        rank: 8,
        alpha: 16.0,
        dropout: 0.0,
        target_modules: vec![],
        merge_weights: false,
    };

    // Test different initialization methods
    let mut lora = LoRALayer::new(256, 256, config, &device).unwrap();

    // Kaiming initialization (default)
    lora.initialize_kaiming().unwrap();
    let params_kaiming = lora.get_lora_a().to_vec().unwrap();

    // Xavier initialization
    lora.initialize_xavier().unwrap();
    let params_xavier = lora.get_lora_a().to_vec().unwrap();

    // Zero initialization for lora_B
    lora.initialize_zero_b().unwrap();
    let params_b = lora.get_lora_b().to_vec().unwrap();

    // Check that B is zeros
    for val in params_b {
        assert_abs_diff_eq!(val, 0.0, epsilon = 1e-6);
    }

    // A matrices should be different between methods
    let mut differences = 0;
    for i in 0..params_kaiming.len() {
        if (params_kaiming[i] - params_xavier[i]).abs() > 1e-6 {
            differences += 1;
        }
    }
    assert!(differences > params_kaiming.len() / 2); // Most values should differ
}