rustytorch/crates/models/rtx-multimodal/tests/video_tests.rs

use approx::assert_abs_diff_eq;
use rtx_multimodal::error::Result;
use rtx_multimodal::video::timesformer::*;
use rtx_tensor::{Device, Tensor};

#[test]
fn test_video_patch_embedding() -> Result<()> {
    let device = Device::default();

    let patch_embed = VideoPatchEmbedding::new(
        224, // image_size
        16,  // patch_size
        3,   // in_channels
        768, // embed_dim
        8,   // num_frames
        &device,
    )?;

    // Should create 14x14 = 196 spatial patches per frame
    assert_eq!(patch_embed.num_patches(), 196);
    assert_eq!(patch_embed.embed_dim(), 768);

    Ok(())
}

#[test]
fn test_video_patch_embedding_forward() -> Result<()> {
    let device = Device::default();
    let patch_embed = VideoPatchEmbedding::new(224, 16, 3, 768, 8, &device)?;

    // Input: [batch_size, channels, frames, height, width]
    let input = Tensor::randn(&[2, 3, 8, 224, 224], &device)?;
    let output = patch_embed.forward(&input)?;

    // Expected: [batch_size, num_frames, num_patches, embed_dim]
    assert_eq!(output.shape(), &[2, 8, 196, 768]);

    Ok(())
}

#[test]
fn test_temporal_positional_encoding() -> Result<()> {
    let device = Device::default();
    let temporal_pe = TemporalPositionalEncoding::new(8, 768, &device)?;

    let embeddings = temporal_pe.forward(8)?;
    // Should return [1, num_frames, embed_dim]
    assert_eq!(embeddings.shape(), &[1, 8, 768]);

    // Test different frame counts
    let shorter = temporal_pe.forward(4)?;
    assert_eq!(shorter.shape(), &[1, 4, 768]);

    Ok(())
}

#[test]
fn test_spatial_positional_encoding() -> Result<()> {
    let device = Device::default();
    let spatial_pe = SpatialPositionalEncoding::new(196, 768, &device)?; // 14x14 patches

    let embeddings = spatial_pe.forward(196)?;
    assert_eq!(embeddings.shape(), &[1, 196, 768]);

    Ok(())
}

#[test]
fn test_divided_space_time_attention() -> Result<()> {
    let device = Device::default();
    let config = DividedSpaceTimeAttentionConfig {
        embed_dim: 768,
        num_heads: 12,
        dropout: 0.1,
    };

    let attention = DividedSpaceTimeAttention::new(&config, &device)?;

    // Input: [batch, frames, patches, embed_dim]
    let input = Tensor::randn(&[2, 8, 196, 768], &device)?;
    let output = attention.forward(&input)?;

    assert_eq!(output.shape(), &[2, 8, 196, 768]);

    Ok(())
}

#[test]
fn test_timesformer_block() -> Result<()> {
    let device = Device::default();
    let config = TimeSformerBlockConfig {
        embed_dim: 768,
        num_heads: 12,
        mlp_ratio: 4.0,
        dropout: 0.1,
        attention_dropout: 0.1,
    };

    let block = TimeSformerBlock::new(&config, &device)?;

    let input = Tensor::randn(&[2, 8, 197, 768], &device)?; // +1 for cls token
    let output = block.forward(&input)?;

    assert_eq!(output.shape(), &[2, 8, 197, 768]);

    Ok(())
}

#[test]
fn test_timesformer_model() -> Result<()> {
    let device = Device::default();
    let config = TimeSformerConfig {
        image_size: 224,
        patch_size: 16,
        in_channels: 3,
        embed_dim: 768,
        depth: 12,
        num_heads: 12,
        mlp_ratio: 4.0,
        num_frames: 8,
        num_classes: 400, // Kinetics-400
        dropout: 0.1,
        attention_dropout: 0.1,
    };

    let model = TimeSformerModel::new(&config, &device)?;

    // Input video: [batch, channels, frames, height, width]
    let input = Tensor::randn(&[2, 3, 8, 224, 224], &device)?;
    let output = model.forward(&input)?;

    // Should output class logits
    assert_eq!(output.shape(), &[2, 400]);

    Ok(())
}

#[test]
fn test_timesformer_feature_extraction() -> Result<()> {
    let device = Device::default();
    let config = TimeSformerConfig {
        image_size: 224,
        patch_size: 16,
        in_channels: 3,
        embed_dim: 768,
        depth: 12,
        num_heads: 12,
        mlp_ratio: 4.0,
        num_frames: 8,
        num_classes: 400,
        dropout: 0.0, // Disable for deterministic testing
        attention_dropout: 0.0,
    };

    let model = TimeSformerModel::new(&config, &device)?;

    let input = Tensor::randn(&[1, 3, 8, 224, 224], &device)?;
    let features = model.extract_features(&input)?;

    // Should return cls token features across time
    assert_eq!(features.shape(), &[1, 8, 768]);

    Ok(())
}

#[test]
fn test_video_frame_sampling() -> Result<()> {
    let device = Device::default();
    let sampler = VideoFrameSampler::new(8, true, &device)?; // 8 frames, uniform sampling

    // Simulate video with 32 frames
    let video = Tensor::randn(&[1, 3, 32, 224, 224], &device)?;
    let sampled = sampler.sample_frames(&video)?;

    // Should sample 8 frames uniformly
    assert_eq!(sampled.shape(), &[1, 3, 8, 224, 224]);

    Ok(())
}

#[test]
fn test_video_augmentation() -> Result<()> {
    let device = Device::default();
    let augmentor = VideoAugmentor::new(
        0.2,  // temporal_crop_ratio
        0.8,  // spatial_crop_ratio
        true, // random_horizontal_flip
        &device,
    )?;

    let input = Tensor::randn(&[2, 3, 8, 224, 224], &device)?;
    let augmented = augmentor.forward(&input)?;

    // Output should have same batch and channel dims
    assert_eq!(augmented.shape()[0], 2); // batch
    assert_eq!(augmented.shape()[1], 3); // channels
    // Frames and spatial dims may be different due to cropping

    Ok(())
}

#[test]
fn test_video_attention_patterns() -> Result<()> {
    let device = Device::default();

    // Test space-only attention (within each frame)
    let space_attn = SpaceOnlyAttention::new(768, 12, 0.1, &device)?;
    let input = Tensor::randn(&[2, 8, 196, 768], &device)?;
    let space_output = space_attn.forward(&input)?;
    assert_eq!(space_output.shape(), &[2, 8, 196, 768]);

    // Test time-only attention (across frames for each patch)
    let time_attn = TimeOnlyAttention::new(768, 12, 0.1, &device)?;
    let time_output = time_attn.forward(&input)?;
    assert_eq!(time_output.shape(), &[2, 8, 196, 768]);

    Ok(())
}

#[test]
fn test_joint_space_time_attention() -> Result<()> {
    let device = Device::default();
    let joint_attn = JointSpaceTimeAttention::new(768, 12, 0.1, &device)?;

    // Input: [batch, frames, patches, embed_dim]
    let input = Tensor::randn(&[2, 8, 196, 768], &device)?;
    let output = joint_attn.forward(&input)?;

    assert_eq!(output.shape(), &[2, 8, 196, 768]);

    Ok(())
}

#[test]
fn test_video_preprocessing() -> Result<()> {
    let device = Device::default();
    let preprocessor = VideoPreprocessor::new(
        224,  // target_size
        8,    // num_frames
        30.0, // fps
        &device,
    )?;

    // Simulate raw video frames (small example)
    let raw_frames: Vec<Vec<Vec<Vec<u8>>>> = vec![
    vec![vec![vec![255u8; 224]; 224]; 3]; 16 // 16 RGB frames of 224x224
    ];

    let processed = preprocessor.preprocess(&raw_frames)?;

    // Should sample to target number of frames
    assert_eq!(processed.shape(), &[1, 3, 8, 224, 224]);

    Ok(())
}