foxhunt/ml/tests/varmap_weight_extraction_test.rs

//! VarMap Weight Extraction Tests (TDD)
//!
//! Tests for extracting real model weights from Candle VarMap for quantization.
//! This replaces stub random weights with actual trained model parameters.

use candle_core::{DType, Device, Tensor};
use candle_nn::{VarBuilder, VarMap};
use std::sync::Arc;

use ml::memory_optimization::quantization::{
    extract_weights_from_varmap, QuantizationConfig, QuantizationType, Quantizer,
};
use ml::MLError;

/// Test 1: Extract single tensor from VarMap (SHOULD FAIL - function doesn't exist yet)
#[test]
fn test_extract_single_tensor_from_varmap() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create a test weight tensor
    let original_weight = Tensor::randn(0.0f32, 1.0f32, (64, 128), &device)?;

    // Insert into VarMap via VarBuilder
    let _weight_var = vs.get_with_hints((64, 128), "layer.weight", candle_nn::Init::Const(0.0))?;

    // Manually set the weight through VarMap data
    let vars_data = varmap.data().lock().unwrap();
    if let Some(var) = vars_data.get("layer.weight") {
        var.set(&original_weight)?;
    }
    drop(vars_data);

    // Extract weight using helper function (THIS WILL FAIL - function doesn't exist)
    let extracted = extract_weights_from_varmap(&varmap, "layer.weight")?;

    // Verify extracted tensor matches original
    let extracted_vec = extracted.flatten_all()?.to_vec1::<f32>()?;
    let original_vec = original_weight.flatten_all()?.to_vec1::<f32>()?;

    assert_eq!(extracted_vec.len(), original_vec.len());
    for (a, b) in extracted_vec.iter().zip(original_vec.iter()) {
        assert!(
            (a - b).abs() < 1e-5,
            "Extracted weight mismatch: {} vs {}",
            a,
            b
        );
    }

    Ok(())
}

/// Test 2: Extract multiple tensors from VarMap
#[test]
fn test_extract_multiple_tensors() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create multiple weight tensors
    let weight1 = Tensor::randn(0.0f32, 1.0f32, (32, 64), &device)?;
    let weight2 = Tensor::randn(0.0f32, 1.0f32, (64, 128), &device)?;
    let bias1 = Tensor::randn(0.0f32, 0.1f32, (64,), &device)?;

    // Insert into VarMap
    let _ = vs.get_with_hints((32, 64), "layer1.weight", candle_nn::Init::Const(0.0))?;
    let _ = vs.get_with_hints((64, 128), "layer2.weight", candle_nn::Init::Const(0.0))?;
    let _ = vs.get_with_hints((64,), "layer1.bias", candle_nn::Init::Const(0.0))?;

    let vars_data = varmap.data().lock().unwrap();
    vars_data.get("layer1.weight").unwrap().set(&weight1)?;
    vars_data.get("layer2.weight").unwrap().set(&weight2)?;
    vars_data.get("layer1.bias").unwrap().set(&bias1)?;
    drop(vars_data);

    // Extract all weights
    let extracted_w1 = extract_weights_from_varmap(&varmap, "layer1.weight")?;
    let extracted_w2 = extract_weights_from_varmap(&varmap, "layer2.weight")?;
    let extracted_b1 = extract_weights_from_varmap(&varmap, "layer1.bias")?;

    // Verify shapes
    assert_eq!(extracted_w1.dims(), &[32, 64]);
    assert_eq!(extracted_w2.dims(), &[64, 128]);
    assert_eq!(extracted_b1.dims(), &[64]);

    Ok(())
}

/// Test 3: Handle missing key error
#[test]
fn test_missing_key_error() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Insert one weight
    let weight = Tensor::randn(0.0f32, 1.0f32, (32, 64), &device)?;
    let _ = vs.get_with_hints((32, 64), "layer.weight", candle_nn::Init::Const(0.0))?;

    let vars_data = varmap.data().lock().unwrap();
    vars_data.get("layer.weight").unwrap().set(&weight)?;
    drop(vars_data);

    // Try to extract non-existent key
    let result = extract_weights_from_varmap(&varmap, "nonexistent.key");

    assert!(result.is_err());
    match result {
        Err(MLError::ModelError(msg)) => {
            assert!(
                msg.contains("not found") || msg.contains("missing"),
                "Expected 'not found' error, got: {}",
                msg
            );
        },
        _ => panic!("Expected ModelError for missing key"),
    }

    Ok(())
}

/// Test 4: Handle dtype preservation (F64 vs F32)
#[test]
fn test_dtype_preservation() -> anyhow::Result<()> {
    let device = Device::Cpu;

    // Test F32
    let varmap_f32 = Arc::new(VarMap::new());
    let vs_f32 = VarBuilder::from_varmap(&varmap_f32, DType::F32, &device);
    let weight_f32 = Tensor::randn(0.0f32, 1.0f32, (10, 20), &device)?;
    let _ = vs_f32.get_with_hints((10, 20), "weight", candle_nn::Init::Const(0.0))?;
    varmap_f32
        .data()
        .lock()
        .unwrap()
        .get("weight")
        .unwrap()
        .set(&weight_f32)?;

    let extracted_f32 = extract_weights_from_varmap(&varmap_f32, "weight")?;
    assert_eq!(extracted_f32.dtype(), DType::F32);

    // Test F64
    let varmap_f64 = Arc::new(VarMap::new());
    let vs_f64 = VarBuilder::from_varmap(&varmap_f64, DType::F64, &device);
    let weight_f64 = Tensor::randn(0.0f64, 1.0f64, (10, 20), &device)?;
    let _ = vs_f64.get_with_hints((10, 20), "weight", candle_nn::Init::Const(0.0))?;
    varmap_f64
        .data()
        .lock()
        .unwrap()
        .get("weight")
        .unwrap()
        .set(&weight_f64)?;

    let extracted_f64 = extract_weights_from_varmap(&varmap_f64, "weight")?;
    assert_eq!(extracted_f64.dtype(), DType::F64);

    Ok(())
}

/// Test 5: Extract from nested VarMap structure (e.g., "encoder.layer1.weight")
#[test]
fn test_nested_key_extraction() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create nested structure
    let weight = Tensor::randn(0.0f32, 1.0f32, (128, 256), &device)?;
    let _ = vs.get_with_hints(
        (128, 256),
        "encoder.layer1.weight",
        candle_nn::Init::Const(0.0),
    )?;
    varmap
        .data()
        .lock()
        .unwrap()
        .get("encoder.layer1.weight")
        .unwrap()
        .set(&weight)?;

    // Extract using nested key
    let extracted = extract_weights_from_varmap(&varmap, "encoder.layer1.weight")?;
    assert_eq!(extracted.dims(), &[128, 256]);

    Ok(())
}

/// Test 6: Quantize using extracted weights (integration test)
#[test]
fn test_quantize_with_extracted_weights() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create model weight with known values
    let weight_data: Vec<f32> = (0..64).map(|i| i as f32 * 0.1).collect();
    let weight = Tensor::from_vec(weight_data.clone(), (8, 8), &device)?;
    let _ = vs.get_with_hints((8, 8), "fc.weight", candle_nn::Init::Const(0.0))?;
    varmap
        .data()
        .lock()
        .unwrap()
        .get("fc.weight")
        .unwrap()
        .set(&weight)?;

    // Extract weight
    let extracted = extract_weights_from_varmap(&varmap, "fc.weight")?;

    // Quantize using Quantizer
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());
    let quantized = quantizer.quantize_tensor(&extracted, "fc.weight")?;

    // Verify quantization succeeded
    assert_eq!(quantized.quant_type, QuantizationType::Int8);
    assert!(quantized.scale > 0.0);

    // Dequantize and verify approximate reconstruction
    let dequantized = quantizer.dequantize_tensor(&quantized)?;
    let dequant_vec = dequantized.flatten_all()?.to_vec1::<f32>()?;

    // Should be close to original (within quantization error)
    for (orig, dequant) in weight_data.iter().zip(dequant_vec.iter()) {
        let error = (orig - dequant).abs();
        assert!(
            error < 0.5,
            "Quantization error too large: {} vs {} (error: {})",
            orig,
            dequant,
            error
        );
    }

    Ok(())
}

/// Test 7: Handle empty VarMap
#[test]
fn test_empty_varmap() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let _vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Try to extract from empty VarMap
    let result = extract_weights_from_varmap(&varmap, "any.key");
    assert!(result.is_err());

    Ok(())
}

/// Test 8: Large tensor extraction (stress test)
#[test]
fn test_large_tensor_extraction() -> anyhow::Result<()> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create large weight tensor (simulating TFT/Transformer layer)
    let large_weight = Tensor::randn(0.0f32, 1.0f32, (1024, 2048), &device)?;
    let _ = vs.get_with_hints(
        (1024, 2048),
        "transformer.layer.weight",
        candle_nn::Init::Const(0.0),
    )?;
    varmap
        .data()
        .lock()
        .unwrap()
        .get("transformer.layer.weight")
        .unwrap()
        .set(&large_weight)?;

    // Extract and verify
    let extracted = extract_weights_from_varmap(&varmap, "transformer.layer.weight")?;
    assert_eq!(extracted.dims(), &[1024, 2048]);
    assert_eq!(extracted.elem_count(), 1024 * 2048);

    Ok(())
}