foxhunt/ml/tests/tft_varmap_checkpoint_test.rs

//! TFT VarMap Checkpoint Serialization Test
//!
//! **Wave 8.5: Validate TFT VarMap Checkpoint Serialization (File-Based Pattern)**
//!
//! Tests that verify the TFT checkpoint save/load works correctly using the
//! file-based VarMap serialization pattern implemented in Wave 6.6.
//!
//! ## File-Based Serialization Pattern (Lines 696-716)
//!
//! ```rust
//! async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
//!     // 1. Create temporary file with UUID
//!     let temp_path = temp_dir.join(format!("tft_checkpoint_{}.safetensors", Uuid::new_v4()));
//!
//!     // 2. Save VarMap to file
//!     self.varmap.save(temp_path_str)?;
//!
//!     // 3. Read file into bytes
//!     let buffer = std::fs::read(&temp_path)?;
//!
//!     // 4. Clean up temp file
//!     let _ = std::fs::remove_file(&temp_path);
//!
//!     Ok(buffer)
//! }
//! ```
//!
//! ## Deserialization Pattern (Lines 718-741)
//!
//! ```rust
//! async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
//!     // 1. Write bytes to temporary file
//!     let temp_path = temp_dir.join(format!("tft_restore_{}.safetensors", Uuid::new_v4()));
//!     std::fs::write(&temp_path, data)?;
//!
//!     // 2. Get mutable access to VarMap (requires Arc::get_mut)
//!     let varmap_mut = Arc::get_mut(&mut self.varmap)?;
//!
//!     // 3. Load checkpoint into VarMap
//!     varmap_mut.load(temp_path_str)?;
//!
//!     // 4. Clean up temp file
//!     let _ = std::fs::remove_file(&temp_path);
//!
//!     Ok(())
//! }
//! ```
//!
//! ## Test Coverage
//!
//! 1. **Basic Save/Load Cycle**: Verify checkpoint saves and loads correctly
//! 2. **State Preservation**: Verify model parameters are restored exactly
//! 3. **Temporary File Cleanup**: Ensure no temp files leak
//! 4. **Concurrent Checkpointing**: Multiple saves should not conflict (UUID isolation)
//! 5. **File Descriptor Management**: No FD leaks after repeated save/load
//! 6. **Arc::get_mut Validation**: Proper mutable access to VarMap
//! 7. **Large Model Checkpoint**: Test with realistic model size
//!
#![allow(unused_crate_dependencies)]

use anyhow::Result;
use ml::checkpoint::{CheckpointConfig, CheckpointManager, Checkpointable};
use ml::tft::{TFTConfig, TemporalFusionTransformer};
use ndarray::{Array1, Array2};
use std::path::PathBuf;
use std::sync::Arc;

/// Test 1: Basic checkpoint save/load cycle
#[tokio::test]
async fn test_tft_varmap_basic_save_load() -> Result<()> {
    println!("\n=== Test 1: Basic VarMap Save/Load ===");

    let checkpoint_dir = PathBuf::from("/tmp/tft_varmap_test_basic");
    std::fs::create_dir_all(&checkpoint_dir)?;

    // Create TFT model
    let config = TFTConfig {
        input_dim: 32,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 3,
        num_static_features: 3,
        num_known_features: 5,
        num_unknown_features: 24  // 3 + 5 + 24 = 32 (fixed feature count mismatch),
        ..Default::default(),
    };

    let mut model = TemporalFusionTransformer::new(config.clone())?;
    model.is_trained = true;

    println!(
        "✓ TFT model created: hidden_dim={}, num_heads={}",
        config.hidden_dim, config.num_heads
    );

    // Create checkpoint manager
    let checkpoint_config = CheckpointConfig {
        base_dir: checkpoint_dir.clone(),
        ..Default::default()
    };
    let manager = CheckpointManager::new(checkpoint_config)?;

    // Save checkpoint (uses serialize_state internally)
    println!("Saving checkpoint...");
    let checkpoint_id = manager.save_checkpoint(&model, None).await?;
    println!("✓ Checkpoint saved: {}", checkpoint_id);

    // Create new model instance
    let mut restored_model = TemporalFusionTransformer::new(config)?;

    // Load checkpoint (uses deserialize_state internally)
    println!("Loading checkpoint...");
    let metadata = manager
        .load_checkpoint(&mut restored_model, &checkpoint_id)
        .await?;
    println!(
        "✓ Checkpoint loaded: epoch={:?}, step={:?}",
        metadata.epoch, metadata.step
    );

    // Verify configuration matches
    assert_eq!(restored_model.config.hidden_dim, model.config.hidden_dim);
    assert_eq!(restored_model.config.num_heads, model.config.num_heads);
    assert_eq!(restored_model.config.num_layers, model.config.num_layers);
    println!("✓ All configuration parameters match");

    Ok(())
}

/// Test 2: State preservation - verify model parameters are restored exactly
#[tokio::test]
async fn test_tft_varmap_state_preservation() -> Result<()> {
    println!("\n=== Test 2: State Preservation ===");

    let checkpoint_dir = PathBuf::from("/tmp/tft_varmap_test_state");
    std::fs::create_dir_all(&checkpoint_dir)?;

    let config = TFTConfig {
        input_dim: 16,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 3,
        num_static_features: 2,
        num_known_features: 4,
        num_unknown_features: 10  // 2 + 4 + 10 = 16 (fixed feature count mismatch),
        ..Default::default(),
    };

    let mut original_model = TemporalFusionTransformer::new(config.clone())?;
    original_model.is_trained = true;

    // Run a prediction to initialize state
    let static_features = Array1::from_vec(vec![1.0, 2.0]);
    let historical_features = Array2::from_shape_vec((20, 8), vec![0.5; 160])?;
    let future_features = Array2::from_shape_vec((5, 4), vec![1.0; 20])?;

    let original_prediction = original_model.predict_horizons(
        &static_features,
        &historical_features,
        &future_features,
    )?;

    println!(
        "✓ Original prediction: {} horizons, latency={}μs",
        original_prediction.predictions.len(),
        original_prediction.latency_us
    );

    // Save checkpoint
    let checkpoint_config = CheckpointConfig {
        base_dir: checkpoint_dir.clone(),
        ..Default::default()
    };
    let manager = CheckpointManager::new(checkpoint_config)?;
    let checkpoint_id = manager.save_checkpoint(&original_model, None).await?;
    println!("✓ Checkpoint saved: {}", checkpoint_id);

    // Load into new model
    let mut restored_model = TemporalFusionTransformer::new(config)?;
    restored_model.is_trained = true;
    manager
        .load_checkpoint(&mut restored_model, &checkpoint_id)
        .await?;
    println!("✓ Checkpoint loaded into new model");

    // Run same prediction on restored model
    let restored_prediction = restored_model.predict_horizons(
        &static_features,
        &historical_features,
        &future_features,
    )?;

    println!(
        "✓ Restored prediction: {} horizons, latency={}μs",
        restored_prediction.predictions.len(),
        restored_prediction.latency_us
    );

    // Verify predictions match (within floating point tolerance)
    assert_eq!(
        original_prediction.predictions.len(),
        restored_prediction.predictions.len()
    );

    for (i, (orig, restored)) in original_prediction
        .predictions
        .iter()
        .zip(restored_prediction.predictions.iter())
        .enumerate()
    {
        let diff = (orig - restored).abs();
        assert!(
            diff < 1e-5,
            "Prediction {} mismatch: original={}, restored={}, diff={}",
            i,
            orig,
            restored,
            diff
        );
    }
    println!("✓ All predictions match within tolerance (1e-5)");

    Ok(())
}

/// Test 3: Temporary file cleanup verification
#[tokio::test]
async fn test_tft_varmap_temp_file_cleanup() -> Result<()> {
    println!("\n=== Test 3: Temporary File Cleanup ===");

    let config = TFTConfig {
        input_dim: 16,
        hidden_dim: 32,
        num_heads: 4,
        ..Default::default()
    };

    let model = TemporalFusionTransformer::new(config)?;

    // Get initial temp file count
    let temp_dir = std::env::temp_dir();
    let initial_files: Vec<_> = std::fs::read_dir(&temp_dir)?
        .filter_map(|entry| entry.ok())
        .filter(|entry| {
            entry
                .file_name()
                .to_string_lossy()
                .starts_with("tft_checkpoint_")
                || entry
                    .file_name()
                    .to_string_lossy()
                    .starts_with("tft_restore_")
        })
        .collect();

    println!("Initial temp files: {}", initial_files.len());

    // Perform multiple save operations
    for i in 0..10 {
        let data = model.serialize_state().await?;
        println!("  Save {}: {} bytes", i + 1, data.len());
    }

    // Check for leaked temp files
    let final_files: Vec<_> = std::fs::read_dir(&temp_dir)?
        .filter_map(|entry| entry.ok())
        .filter(|entry| {
            entry
                .file_name()
                .to_string_lossy()
                .starts_with("tft_checkpoint_")
                || entry
                    .file_name()
                    .to_string_lossy()
                    .starts_with("tft_restore_")
        })
        .collect();

    println!("Final temp files: {}", final_files.len());

    // Should have same number of temp files (cleanup working)
    assert_eq!(
        initial_files.len(),
        final_files.len(),
        "Temporary files leaked: initial={}, final={}",
        initial_files.len(),
        final_files.len()
    );

    println!("✓ No temporary file leaks detected");

    Ok(())
}

/// Test 4: Concurrent checkpointing - UUID isolation prevents conflicts
#[tokio::test]
async fn test_tft_varmap_concurrent_saves() -> Result<()> {
    println!("\n=== Test 4: Concurrent Checkpointing ===");

    let config = TFTConfig {
        input_dim: 16,
        hidden_dim: 32,
        num_heads: 4,
        ..Default::default()
    };

    // Create multiple models
    let models: Vec<_> = (0..5)
        .map(|_| TemporalFusionTransformer::new(config.clone()))
        .collect::<Result<Vec<_>, _>>()?;

    println!(
        "Created {} TFT models for concurrent save test",
        models.len()
    );

    // Save all models concurrently
    let save_tasks: Vec<_> = models
        .iter()
        .enumerate()
        .map(|(i, model)| {
            let model_ref = model;
            async move {
                let data = model_ref.serialize_state().await?;
                println!("  Model {} saved: {} bytes", i, data.len());
                Ok::<_, anyhow::Error>(data)
            }
        })
        .collect();

    let results = futures::future::try_join_all(save_tasks).await?;

    assert_eq!(results.len(), 5, "All 5 models should save successfully");
    println!(
        "✓ All {} concurrent saves completed without conflicts",
        results.len()
    );

    // Verify all saved data is non-empty and different
    for (i, data) in results.iter().enumerate() {
        assert!(!data.is_empty(), "Model {} data should not be empty", i);
    }
    println!("✓ All saved data is valid");

    Ok(())
}

/// Test 5: File descriptor leak check
#[tokio::test]
async fn test_tft_varmap_fd_leak() -> Result<()> {
    println!("\n=== Test 5: File Descriptor Leak Check ===");

    let config = TFTConfig {
        input_dim: 16,
        hidden_dim: 32,
        num_heads: 4,
        ..Default::default()
    };

    let mut model = TemporalFusionTransformer::new(config.clone())?;

    // Get process FD count (Linux-specific)
    let get_fd_count = || -> Result<usize> {
        let pid = std::process::id();
        let fd_dir = format!("/proc/{}/fd", pid);
        Ok(std::fs::read_dir(&fd_dir)?.count())
    };

    let initial_fds = get_fd_count().unwrap_or(0);
    println!("Initial FD count: {}", initial_fds);

    // Perform 50 save/load cycles
    for i in 0..50 {
        // Save
        let data = model.serialize_state().await?;

        // Load
        model.deserialize_state(&data).await?;

        if (i + 1) % 10 == 0 {
            println!("  Completed {} save/load cycles", i + 1);
        }
    }

    let final_fds = get_fd_count().unwrap_or(0);
    println!("Final FD count: {}", final_fds);

    // Allow small variance (tokio runtime may open/close FDs)
    let fd_diff = (final_fds as i32 - initial_fds as i32).abs();
    assert!(
        fd_diff < 10,
        "Significant FD leak detected: initial={}, final={}, diff={}",
        initial_fds,
        final_fds,
        fd_diff
    );

    println!("✓ No file descriptor leaks detected (diff={})", fd_diff);

    Ok(())
}

/// Test 6: Arc::get_mut validation - proper mutable access
#[tokio::test]
async fn test_tft_varmap_arc_get_mut() -> Result<()> {
    println!("\n=== Test 6: Arc::get_mut Validation ===");

    let config = TFTConfig {
        input_dim: 16,
        hidden_dim: 32,
        num_heads: 4,
        ..Default::default()
    };

    let mut model = TemporalFusionTransformer::new(config)?;

    // Save state
    let data = model.serialize_state().await?;
    println!("✓ Serialized state: {} bytes", data.len());

    // Load state (requires Arc::get_mut)
    model.deserialize_state(&data).await?;
    println!("✓ Deserialized state successfully (Arc::get_mut worked)");

    // Verify model is still usable
    let test_input_static = vec![1.0f32; 2];
    let test_input_hist = vec![0.5f32; 20 * 8];
    let test_input_fut = vec![1.0f32; 5 * 4];

    let predictions = model.predict_fast(&test_input_static, &test_input_hist, &test_input_fut)?;

    assert_eq!(predictions.len(), 5, "Should predict 5 horizons");
    println!(
        "✓ Model operational after load: {} predictions",
        predictions.len()
    );

    Ok(())
}

/// Test 7: Large model checkpoint - realistic model size
#[tokio::test]
async fn test_tft_varmap_large_model() -> Result<()> {
    println!("\n=== Test 7: Large Model Checkpoint ===");

    let checkpoint_dir = PathBuf::from("/tmp/tft_varmap_test_large");
    std::fs::create_dir_all(&checkpoint_dir)?;

    // Create large TFT model (production-sized)
    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 256,
        num_heads: 16,
        num_layers: 6,
        prediction_horizon: 50,
        sequence_length: 200,
        num_quantiles: 9,
        num_static_features: 10,
        num_known_features: 20,
        num_unknown_features: 34  // 10 + 20 + 34 = 64 (fixed feature count mismatch),
        ..Default::default(),
    };

    let mut model = TemporalFusionTransformer::new(config.clone())?;
    model.is_trained = true;

    println!("✓ Large TFT model created:");
    println!("  - Hidden dim: {}", config.hidden_dim);
    println!("  - Num heads: {}", config.num_heads);
    println!("  - Num layers: {}", config.num_layers);
    println!("  - Prediction horizon: {}", config.prediction_horizon);

    // Measure save time
    let save_start = std::time::Instant::now();
    let data = model.serialize_state().await?;
    let save_time = save_start.elapsed();
    println!("✓ Save time: {:?} ({} bytes)", save_time, data.len());

    // Measure load time
    let load_start = std::time::Instant::now();
    model.deserialize_state(&data).await?;
    let load_time = load_start.elapsed();
    println!("✓ Load time: {:?}", load_time);

    // Verify model works after load
    let checkpoint_config = CheckpointConfig {
        base_dir: checkpoint_dir.clone(),
        ..Default::default()
    };
    let manager = CheckpointManager::new(checkpoint_config)?;
    let checkpoint_id = manager.save_checkpoint(&model, None).await?;
    println!("✓ Full checkpoint saved: {}", checkpoint_id);

    // Load and verify
    let mut restored_model = TemporalFusionTransformer::new(config)?;
    restored_model.is_trained = true;
    manager
        .load_checkpoint(&mut restored_model, &checkpoint_id)
        .await?;
    println!("✓ Large model restored successfully");

    // Performance check
    assert!(
        save_time.as_millis() < 1000,
        "Save time too slow: {:?} (should be <1s)",
        save_time
    );
    assert!(
        load_time.as_millis() < 1000,
        "Load time too slow: {:?} (should be <1s)",
        load_time
    );
    println!("✓ Performance within acceptable limits");

    Ok(())
}

/// Test 8: Multiple save/load cycles - stress test
#[tokio::test]
async fn test_tft_varmap_repeated_cycles() -> Result<()> {
    println!("\n=== Test 8: Repeated Save/Load Cycles ===");

    let config = TFTConfig {
        input_dim: 32,
        hidden_dim: 64,
        num_heads: 4,
        ..Default::default()
    };

    let mut model = TemporalFusionTransformer::new(config)?;

    // Perform 100 save/load cycles
    let num_cycles = 100;
    let mut total_save_time = std::time::Duration::ZERO;
    let mut total_load_time = std::time::Duration::ZERO;

    for i in 0..num_cycles {
        // Save
        let save_start = std::time::Instant::now();
        let data = model.serialize_state().await?;
        total_save_time += save_start.elapsed();

        // Load
        let load_start = std::time::Instant::now();
        model.deserialize_state(&data).await?;
        total_load_time += load_start.elapsed();

        if (i + 1) % 20 == 0 {
            println!("  Completed {} cycles", i + 1);
        }
    }

    let avg_save_time = total_save_time / num_cycles;
    let avg_load_time = total_load_time / num_cycles;

    println!("✓ Completed {} save/load cycles", num_cycles);
    println!("  - Average save time: {:?}", avg_save_time);
    println!("  - Average load time: {:?}", avg_load_time);
    println!("  - Total time: {:?}", total_save_time + total_load_time);

    // Performance assertions
    assert!(
        avg_save_time.as_millis() < 100,
        "Average save too slow: {:?}",
        avg_save_time
    );
    assert!(
        avg_load_time.as_millis() < 100,
        "Average load too slow: {:?}",
        avg_load_time
    );

    println!("✓ All cycles completed within performance targets");

    Ok(())
}