import torch
from torch.nn import Module
from torch.utils.data import Dataset, DataLoader

# mock offline dataset

class MockOfflineRoboticFrameDataset(Module):
    """
    A mock dataset that returns an offline video prefix to seed the imagination
    of the robotics reinforcement learning / trajectory generation pipeline.
    
    Returns a dictionary with:
      - 'video': The seed video prefix [T, C, H, W]
      - 'proprioception': The robot state corresponding to the video [T, Proprio_Dim]
      - 'prompt_token_ids': Integer token IDs for the task instruction
    """
    def __init__(
        self,
        num_samples = 100,
        prefix_frames = 8,
        channels = 3,
        image_size = 64,
        proprio_dim = 14,
        action_dim = 10,
        action_chunk_len = 8,
        num_prompt_tokens = 12,
        dtype = torch.bfloat16
    ):
        super().__init__()
        self.num_samples = num_samples
        self.prefix_frames = prefix_frames
        self.channels = channels
        self.image_size = image_size
        self.proprio_dim = proprio_dim
        self.action_dim = action_dim
        self.action_chunk_len = action_chunk_len
        self.num_prompt_tokens = num_prompt_tokens
        self.dtype = dtype

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # Generate a random mock RGB video prefix
        video = torch.randn(
            self.prefix_frames,
            self.channels,
            self.image_size,
            self.image_size,
            dtype = self.dtype
        )
        
        # Mock "next" video frame for flow-matching loss
        video_next = torch.randn(
            1, # Single next frame
            self.channels,
            self.image_size,
            self.image_size,
            dtype = self.dtype
        )
        
        # Generate mock proprioception
        proprioception = torch.randn(
            self.prefix_frames,
            self.proprio_dim,
            dtype = self.dtype
        )
        
        # Mock actions (e.g., velocity commands)
        # Total sequence must match prefix + future (imagination)
        # In this mock, we assume 1 future frame is predicted by default
        actions = torch.randn(
            self.prefix_frames + 1,
            10,
            dtype = self.dtype
        )
        
        # Mock token IDs
        prompt_token_ids = torch.randint(0, 1000, (self.num_prompt_tokens,))
        
        return {
            "video": video,
            "next_video": video_next,
            "actions": actions,
            "proprioception": proprioception,
            "prompt_token_ids": prompt_token_ids
        }
