import torch
from mimic_video.cosmos_predict import CosmosPredictWrapper

def test_rtc_forward():
    print("Testing forward with fixed video prefix...")
    wrapper = CosmosPredictWrapper(tiny = True, random_weights = True, train_fixed_video_prefix = True, train_fixed_video_prefix_max_delay = 2)
    wrapper.train()
    
    batch = 2
    frames = 4
    videos = torch.randn(batch, frames, 3, 64, 64)
    prompts = ["a video of a robot", "a video of a cat"]
    
    # forward should work without error
    hiddens = wrapper(videos, prompts = prompts)
    print("Forward successful.")
    assert hiddens is not None

def test_rtc_finetune():
    print("Testing finetune with fixed video prefix...")
    wrapper = CosmosPredictWrapper(tiny = True, random_weights = True)
    
    class DummyDataset(torch.utils.data.Dataset):
        def __init__(self, n=4):
            self.n = n
        def __len__(self):
            return self.n
        def __getitem__(self, idx):
            return torch.randn(4, 3, 64, 64), "dummy prompt"

    dataset = DummyDataset()
    
    # finetune should work with train_fixed_video_prefix_max_delay
    wrapper.finetune(
        dataset,
        epochs = 1,
        batch_size = 2,
        train_fixed_video_prefix_max_delay = 2
    )
    print("Finetune successful.")

def test_rtc_disabled():
    print("Testing forward and finetune with fixed video prefix disabled...")
    # forward
    wrapper = CosmosPredictWrapper(tiny = True, random_weights = True, train_fixed_video_prefix = False)
    wrapper.train()
    videos = torch.randn(2, 4, 3, 64, 64)
    hiddens = wrapper(videos, prompts = ["a", "b"])
    assert hiddens is not None
    print("Forward (disabled) successful.")

    # finetune
    class DummyDataset(torch.utils.data.Dataset):
        def __init__(self, n=2): self.n = n
        def __len__(self): return self.n
        def __getitem__(self, idx): return torch.randn(4, 3, 64, 64), "dummy"

    wrapper.finetune(DummyDataset(), epochs = 1, batch_size = 2, train_fixed_video_prefix_max_delay = 0)
    print("Finetune (disabled) successful.")

if __name__ == "__main__":
    try:
        test_rtc_forward()
        test_rtc_finetune()
        test_rtc_disabled()
        print("All tests passed!")
    except Exception as e:
        print(f"Test failed with error: {e}")
        import traceback
        traceback.print_exc()
