import torch
import torch.nn as nn
import torch.nn.functional as F

class TextEncoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(embed_dim, nhead=8),
            num_layers=6
        )
        
    def forward(self, text):
        x = self.embedding(text)
        return self.transformer(x)

class VideoGenerator(nn.Module):
    def __init__(self, latent_dim, num_frames, frame_size):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_frames = num_frames
        
        self.generator = nn.Sequential(
            nn.ConvTranspose3d(latent_dim, 512, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm3d(512),
            nn.ReLU(),
            nn.ConvTranspose3d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm3d(256),
            nn.ReLU(),
            nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm3d(128),
            nn.ReLU(),
            nn.ConvTranspose3d(128, 3, kernel_size=4, stride=2, padding=1),
            nn.Tanh()
        )
    
    def forward(self, z):
        return self.generator(z)

class Text2VideoModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, latent_dim, num_frames, frame_size):
        super().__init__()
        self.text_encoder = TextEncoder(vocab_size, embed_dim, hidden_dim=512)
        self.video_generator = VideoGenerator(latent_dim, num_frames, frame_size)
        self.latent_mapper = nn.Linear(embed_dim, latent_dim * num_frames)
        
    def forward(self, text):
        text_features = self.text_encoder(text)
        latent_vector = self.latent_mapper(text_features.mean(dim=1))
        latent_video = latent_vector.view(-1, self.video_generator.latent_dim, 1, 1, 1)
        generated_video = self.video_generator(latent_video)
        return generated_video