"""
Test CosmosPredictWrapper with denoising steps on CPU
"""

import torch
import time

def test_cosmos_denoising():
    print("Testing CosmosPredictWrapper with denoising steps...")
    
    from mimic_video import CosmosPredictWrapper
    
    start = time.time()
    wrapper = CosmosPredictWrapper(
        device = "cpu",
        torch_dtype = torch.float32,
        extract_layer = 1,
        random_weights = True,
        tiny = True
    )
    print(f"Init took: {time.time() - start:.2f}s")
    print(f"Wrapper dim_latent_hidden: {wrapper.dim_latent_hidden}")
    
    # Create a tiny dummy video
    batch_size = 1
    num_frames = 5
    video = torch.rand(batch_size, num_frames, 3, 32, 32)
    
    print(f"Input video shape: {video.shape}")
    
    # Test with 1 step
    print("\n--- 1 denoising step ---")
    start = time.time()
    hidden_1 = wrapper(videos=video, prompts="test", num_inference_steps=1)
    print(f"Forward took: {time.time() - start:.2f}s")
    print(f"Hidden states shape: {hidden_1.shape}")
    
    # Test with 4 steps
    print("\n--- 4 denoising steps ---")
    start = time.time()
    hidden_4 = wrapper(videos=video, prompts="test", num_inference_steps=4)
    print(f"Forward took: {time.time() - start:.2f}s")
    print(f"Hidden states shape: {hidden_4.shape}")
    print(f"Number of cached states: {len(wrapper.cached_hidden_states)}")
    
    assert hidden_1.shape[-1] == wrapper.dim_latent_hidden
    assert hidden_4.shape[-1] == wrapper.dim_latent_hidden
    
    print("\n✓ Denoising test passed!")

if __name__ == "__main__":
    test_cosmos_denoising()
