foxhunt/ml/tests/checkpoint_integrity.rs

//! Generic Checkpoint Integrity Tests
//!
//! This test suite validates checkpoint saving/loading for all ML models
//! to catch VarMap registration bugs where layers are not properly saved.
//!
//! ## Tests Included
//! 1. Parameter Count Validation - Ensures all parameters are saved
//! 2. Checkpoint Restore - Ensures loaded weights match original
//! 3. Layer-by-Layer Parameter Test - Verifies each layer is in checkpoint
//! 4. Checkpoint Size Validation - Ensures checkpoint is reasonable size
//!
//! ## Bug Context
//! MAMBA-2 had critical bug: 90% of model not saved due to SSD layers
//! creating local VarMap instead of using parent VarMap. These tests
//! would have caught it immediately.

use candle_core::{Device, Tensor, DType};
use ml::mamba::{Mamba2Config, Mamba2SSM};
use ml::MLError;
use tempfile::TempDir;
use std::path::PathBuf;

// ============================================================================
// TEST UTILITIES
// ============================================================================

/// Create a temporary directory for test artifacts
fn create_temp_dir() -> TempDir {
    TempDir::new().expect("Failed to create temp directory")
}

/// Get checkpoint path in temp directory
fn checkpoint_path(temp_dir: &TempDir, filename: &str) -> PathBuf {
    temp_dir.path().join(filename)
}

/// Count parameters in a safetensors file
fn count_checkpoint_parameters(path: &PathBuf) -> Result<usize, MLError> {
    use std::collections::HashMap;

    let tensors: HashMap<String, Tensor> = candle_core::safetensors::load(path, &Device::Cpu)
        .map_err(|e| MLError::CheckpointError(format!("Failed to load checkpoint: {}", e)))?;

    let mut total_params = 0;
    for (_name, tensor) in tensors.iter() {
        let shape = tensor.shape();
        let param_count: usize = shape.dims().iter().product();
        total_params += param_count;
    }

    Ok(total_params)
}

/// Count expected parameters from model architecture
fn count_expected_mamba2_parameters(config: &Mamba2Config) -> usize {
    let d_inner = config.d_model * config.expand;

    // Input projection: d_model -> d_inner
    let input_proj = config.d_model * d_inner + d_inner; // weights + bias

    // Output projection: d_inner -> 1 (regression)
    let output_proj = d_inner * 1 + 1; // weights + bias

    // Per-layer parameters
    let mut layer_params = 0;
    for _ in 0..config.num_layers {
        // Layer norm: d_inner (weight + bias)
        layer_params += d_inner * 2;

        // SSD layer projections (THIS IS WHAT WAS MISSING IN CHECKPOINTS)
        // QKV projection: d_model -> 3 * d_head * num_heads
        let qkv_dim = 3 * config.d_head * config.num_heads;
        layer_params += config.d_model * qkv_dim + qkv_dim; // weights + bias

        // Output projection: d_head * num_heads -> d_model
        let out_dim = config.d_head * config.num_heads;
        layer_params += out_dim * config.d_model + config.d_model; // weights + bias

        // State projection: d_model -> d_state
        layer_params += config.d_model * config.d_state + config.d_state; // weights + bias

        // Gate projection: d_model -> d_model
        layer_params += config.d_model * config.d_model + config.d_model; // weights + bias
    }

    input_proj + output_proj + layer_params
}

// ============================================================================
// MAMBA-2 CHECKPOINT INTEGRITY TESTS
// ============================================================================

#[test]
fn test_mamba2_checkpoint_parameter_count() {
    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0,
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Calculate expected parameter count
    let expected_params = count_expected_mamba2_parameters(&config);
    println!("Expected parameters: {}", expected_params);

    // Save checkpoint
    let temp_dir = create_temp_dir();
    let ckpt_path = checkpoint_path(&temp_dir, "mamba2_param_count.safetensors");

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            model.save_checkpoint(ckpt_path.to_str().unwrap()).await
        })
        .expect("Failed to save checkpoint");

    // Count actual parameters in checkpoint
    let actual_params = count_checkpoint_parameters(&ckpt_path)
        .expect("Failed to count checkpoint parameters");

    println!("Actual parameters in checkpoint: {}", actual_params);

    // CRITICAL TEST: Actual should be within 5% of expected
    // If this fails, it means layers are not being saved (VarMap bug)
    let diff_pct = ((actual_params as f64 - expected_params as f64) / expected_params as f64).abs() * 100.0;

    assert!(
        diff_pct < 5.0,
        "Parameter count mismatch! Expected: {}, Actual: {}, Diff: {:.2}%\n\
         This indicates layers are not properly registered in VarMap.",
        expected_params, actual_params, diff_pct
    );
}

#[test]
fn test_mamba2_checkpoint_restore_determinism() {
    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0, // No dropout for determinism
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Create test input
    let input_data: Vec<f64> = (0..80).map(|i| (i as f64) * 0.01).collect();
    let input = Tensor::from_vec(input_data, (1, 10, 8), &device).expect("Failed to create tensor");

    // Run inference BEFORE saving
    let output1 = model.forward(&input).expect("Failed to run forward pass");
    let output1_vec = output1.flatten_all()
        .expect("Failed to flatten")
        .to_vec1::<f64>()
        .expect("Failed to extract values");

    println!("Output before save: {:?}", &output1_vec[..5]);

    // Save checkpoint
    let temp_dir = create_temp_dir();
    let ckpt_path = checkpoint_path(&temp_dir, "mamba2_restore.safetensors");

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            model.save_checkpoint(ckpt_path.to_str().unwrap()).await
        })
        .expect("Failed to save checkpoint");

    // Create NEW model and load checkpoint
    let mut model2 = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model 2");

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            model2.load_checkpoint(ckpt_path.to_str().unwrap()).await
        })
        .expect("Failed to load checkpoint");

    // Run inference AFTER loading
    let output2 = model2.forward(&input).expect("Failed to run forward pass on loaded model");
    let output2_vec = output2.flatten_all()
        .expect("Failed to flatten")
        .to_vec1::<f64>()
        .expect("Failed to extract values");

    println!("Output after load: {:?}", &output2_vec[..5]);

    // CRITICAL TEST: Outputs should be IDENTICAL (within floating point precision)
    // If this fails, it means weights were not properly restored
    assert_eq!(
        output1_vec.len(),
        output2_vec.len(),
        "Output shapes don't match after checkpoint restore"
    );

    for (i, (val1, val2)) in output1_vec.iter().zip(output2_vec.iter()).enumerate() {
        let diff = (val1 - val2).abs();
        assert!(
            diff < 1e-6,
            "Output mismatch at index {}! Before: {}, After: {}, Diff: {}\n\
             This indicates checkpoint did not restore all weights correctly.",
            i, val1, val2, diff
        );
    }
}

#[test]
fn test_mamba2_all_layers_in_checkpoint() {
    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0,
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Save checkpoint
    let temp_dir = create_temp_dir();
    let ckpt_path = checkpoint_path(&temp_dir, "mamba2_layers.safetensors");

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            model.save_checkpoint(ckpt_path.to_str().unwrap()).await
        })
        .expect("Failed to save checkpoint");

    // Load checkpoint and inspect layer names
    use std::collections::HashMap;
    let tensors: HashMap<String, Tensor> = candle_core::safetensors::load(&ckpt_path, &device)
        .expect("Failed to load checkpoint");

    println!("\nCheckpoint contains {} tensors:", tensors.len());
    for name in tensors.keys() {
        println!("  - {}", name);
    }

    // CRITICAL TEST: Verify each expected layer has parameters

    // Input projection
    assert!(
        tensors.contains_key("input_proj.weight"),
        "Missing input_proj.weight in checkpoint!"
    );

    // Output projection
    assert!(
        tensors.contains_key("output_proj.weight"),
        "Missing output_proj.weight in checkpoint!"
    );

    // Layer norms
    for i in 0..config.num_layers {
        let ln_key = format!("ln_{}.weight", i);
        assert!(
            tensors.contains_key(&ln_key),
            "Missing {} in checkpoint!",
            ln_key
        );
    }

    // SSD layers (THIS IS THE CRITICAL BUG - these were MISSING)
    for i in 0..config.num_layers {
        let ssd_prefix = format!("ssd_layer_{}", i);

        // QKV projection
        let qkv_key = format!("{}.qkv_proj.weight", ssd_prefix);
        assert!(
            tensors.contains_key(&qkv_key),
            "CRITICAL BUG: Missing {} in checkpoint!\n\
             This is the VarMap registration bug - SSD layers not saved.",
            qkv_key
        );

        // Output projection
        let out_key = format!("{}.out_proj.weight", ssd_prefix);
        assert!(
            tensors.contains_key(&out_key),
            "CRITICAL BUG: Missing {} in checkpoint!\n\
             This is the VarMap registration bug - SSD layers not saved.",
            out_key
        );

        // State projection
        let state_key = format!("{}.state_proj.weight", ssd_prefix);
        assert!(
            tensors.contains_key(&state_key),
            "CRITICAL BUG: Missing {} in checkpoint!\n\
             This is the VarMap registration bug - SSD layers not saved.",
            state_key
        );

        // Gate projection
        let gate_key = format!("{}.gate_proj.weight", ssd_prefix);
        assert!(
            tensors.contains_key(&gate_key),
            "CRITICAL BUG: Missing {} in checkpoint!\n\
             This is the VarMap registration bug - SSD layers not saved.",
            gate_key
        );
    }
}

#[test]
fn test_mamba2_checkpoint_size_validation() {
    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0,
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Calculate expected size (F64 = 8 bytes per parameter)
    let expected_params = count_expected_mamba2_parameters(&config);
    let expected_size_bytes = expected_params * 8;
    let expected_size_kb = expected_size_bytes as f64 / 1024.0;

    println!("Expected checkpoint size: {:.2} KB ({} params)", expected_size_kb, expected_params);

    // Save checkpoint
    let temp_dir = create_temp_dir();
    let ckpt_path = checkpoint_path(&temp_dir, "mamba2_size.safetensors");

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            model.save_checkpoint(ckpt_path.to_str().unwrap()).await
        })
        .expect("Failed to save checkpoint");

    // Check actual file size
    let metadata = std::fs::metadata(&ckpt_path).expect("Failed to get file metadata");
    let actual_size_kb = metadata.len() as f64 / 1024.0;

    println!("Actual checkpoint size: {:.2} KB", actual_size_kb);

    // CRITICAL TEST: File size should be reasonable (within 20% of expected)
    // If file is too small, layers are missing (VarMap bug)
    // If file is too large, there's metadata overhead (acceptable)

    let size_ratio = actual_size_kb / expected_size_kb;

    assert!(
        size_ratio > 0.8,
        "Checkpoint file is suspiciously small! Expected: {:.2} KB, Actual: {:.2} KB (ratio: {:.2})\n\
         This indicates layers are not being saved (VarMap bug).",
        expected_size_kb, actual_size_kb, size_ratio
    );

    assert!(
        size_ratio < 2.0,
        "Checkpoint file is unexpectedly large! Expected: {:.2} KB, Actual: {:.2} KB (ratio: {:.2})\n\
         This may indicate duplicate parameters or excessive metadata.",
        expected_size_kb, actual_size_kb, size_ratio
    );
}

#[test]
fn test_mamba2_checkpoint_missing_layers_detection() {
    // This test simulates the BUG scenario where SSD layers create local VarMap
    // It should FAIL until the bug is fixed

    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0,
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    let temp_dir = create_temp_dir();
    let ckpt_path = checkpoint_path(&temp_dir, "mamba2_bug_detection.safetensors");

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            model.save_checkpoint(ckpt_path.to_str().unwrap()).await
        })
        .expect("Failed to save checkpoint");

    // Load checkpoint and count layer-specific tensors
    use std::collections::HashMap;
    let tensors: HashMap<String, Tensor> = candle_core::safetensors::load(&ckpt_path, &device)
        .expect("Failed to load checkpoint");

    // Count input/output projection tensors
    let io_tensors = tensors.keys().filter(|k| k.contains("input_proj") || k.contains("output_proj")).count();

    // Count SSD layer tensors (THE BUG: these should exist but don't)
    let ssd_tensors = tensors.keys().filter(|k| k.contains("ssd_layer_")).count();

    // Count layer norm tensors
    let ln_tensors = tensors.keys().filter(|k| k.contains("ln_")).count();

    println!("\nTensor distribution:");
    println!("  Input/Output projections: {}", io_tensors);
    println!("  SSD layers: {}", ssd_tensors);
    println!("  Layer norms: {}", ln_tensors);
    println!("  Total: {}", tensors.len());

    // CRITICAL TEST: SSD tensors should be the MAJORITY of the checkpoint
    // Expected: 4 projections per SSD layer * 2 layers * 2 tensors (weight+bias) = 16 SSD tensors
    // If ssd_tensors is 0, the VarMap bug exists

    let expected_ssd_tensors = config.num_layers * 4 * 2; // 4 projections, 2 tensors each (weight+bias)

    assert!(
        ssd_tensors >= expected_ssd_tensors,
        "CRITICAL BUG DETECTED: Only {} SSD tensors found, expected at least {}!\n\
         This is the VarMap registration bug - SSD layers are creating local VarMap\n\
         instead of using parent VarMap.",
        ssd_tensors, expected_ssd_tensors
    );
}