"""
Test CosmosPredictWrapper
"""

import torch
from mimic_video import CosmosPredictWrapper

def test_batching():
    print("Testing batching support...")
    batch_size = 2
    wrapper = CosmosPredictWrapper(random_weights = True, tiny = True, extract_layer = 19, device = "cpu")

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

    hidden = wrapper(video, prompts = ['a cat', 'a dog'])
    print(f'hidden_states: {hidden.shape}')

    assert hidden.shape == (batch_size, 192, 16)
    print('batching test success')

def test_single():
    print("\nTesting single video support...")
    wrapper = CosmosPredictWrapper(random_weights = True, tiny = True, extract_layer = 19, device = "cpu")

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

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

    assert hidden.shape == (1, 192, 16)
    print('single test success')

def test_prompt_token_ids():
    print("\nTesting prompt_token_ids support...")
    batch_size = 2
    wrapper = CosmosPredictWrapper(random_weights = True, tiny = True, extract_layer = 19, device = "cpu")

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

    # Mock token ids
    prompt_token_ids = torch.randint(0, 32128, (batch_size, 16))

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

    assert hidden.shape == (batch_size, 192, 16)
    print('prompt_token_ids test success')

if __name__ == '__main__':
    test_batching()
    test_single()
    test_prompt_token_ids()
