"""
Test Cosmos2_5PredictWrapper
"""

import torch
from mimic_video import Cosmos2_5PredictWrapper, MimicVideo

def test_cosmos_2_5_wrapper():
    print("Testing Cosmos2_5PredictWrapper basic functionality...")
    batch_size = 1
    wrapper = Cosmos2_5PredictWrapper(random_weights = True, tiny = True, extract_layer = 1)

    assert isinstance(wrapper, Cosmos2_5PredictWrapper)
    print(f"Verified class: {type(wrapper).__name__}")

    video = torch.rand(batch_size, 5, 3, 64, 64)
    print(f'input: {video.shape}')

    hidden = wrapper(video)
    print(f'hidden_states: {hidden.shape}')

    assert hidden.ndim == 3
    print('Cosmos2_5PredictWrapper test success')

def test_cosmos_2_5_integration():
    print("\nTesting Cosmos2_5PredictWrapper integration with MimicVideo...")
    wrapper = Cosmos2_5PredictWrapper(random_weights = True, tiny = False, extract_layer = 1)
    
    model = MimicVideo(
        dim = 512,
        video_predict_wrapper = wrapper,
        depth = 2
    )

    video = torch.rand(1, 5, 3, 32, 32)
    actions = torch.randn(1, 32, 20)
    joint_state = torch.randn(1, 32)

    loss = model(
        video = video,
        actions = actions,
        joint_state = joint_state,
        prompts = 'test prompt'
    )

    print(f'Loss: {loss.item()}')
    assert loss.numel() == 1
    
    sampled_actions = model.sample(
        video = video,
        joint_state = joint_state,
        prompts = 'sample prompt'
    )
    
    print(f'Sampleed actions shape: {sampled_actions.shape}')
    assert sampled_actions.shape == (1, 32, 20)
    print('MimicVideo integration test success')

if __name__ == '__main__':
    test_cosmos_2_5_wrapper()
    test_cosmos_2_5_integration()
