import torch
import torch.nn as nn
import torch.nn.functional as F
from torchaudio.transforms import MelSpectrogram

class ToyTextEncoder(nn.Module):
    """Processes integer token sequences into a pooled text embedding."""
    def __init__(self, vocab_size=10000, embed_dim=256, out_dim=512):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=embed_dim, nhead=4, batch_first=True),
            num_layers=2
        )
        self.proj = nn.Linear(embed_dim, out_dim)

    def forward(self, x):
        # x: [Batch, Seq_Len]
        emb = self.embedding(x)
        out = self.encoder(emb)
        pooled = out.mean(dim=1)  # Mean pooling over sequence length
        return self.proj(pooled)

class ToyVisionEncoder(nn.Module):
    """Processes (B, C, H, W) images into an image embedding."""
    def __init__(self, out_dim=512):
        super().__init__()
        self.conv_stem = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        self.proj = nn.Linear(64, out_dim)

    def forward(self, x):
        # x: [Batch, 3, H, W]
        feats = self.conv_stem(x).squeeze(-1).squeeze(-1)
        return self.proj(feats)

class ToyAudioEncoder(nn.Module):
    """Transforms raw waveforms -> Mel-Spectrogram -> Audio embedding."""
    def __init__(self, sample_rate=16000, n_mels=64, out_dim=512):
        super().__init__()
        self.spectrogram = MelSpectrogram(sample_rate=sample_rate, n_mels=n_mels)
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        self.proj = nn.Linear(32, out_dim)

    def forward(self, x):
        # x: [Batch, Waveform_Samples]
        with torch.no_grad():
            spec = self.spectrogram(x).unsqueeze(1)  # Shape: [Batch, 1, Mels, Time]
        feats = self.conv(spec).squeeze(-1).squeeze(-1)
        return self.proj(feats)

class UnifiedMultimodalSystem(nn.Module):
    """Projects Text, Vision, and Audio into a single L2-normalized metric space."""
    def __init__(self, shared_dim=512):
        super().__init__()
        self.text_enc = ToyTextEncoder(out_dim=shared_dim)
        self.vis_enc = ToyVisionEncoder(out_dim=shared_dim)
        self.aud_enc = ToyAudioEncoder(out_dim=shared_dim)

    def forward(self, text=None, image=None, audio=None):
        embeddings = {}
        if text is not None:
            embeddings["text"] = F.normalize(self.text_enc(text), p=2, dim=-1)
        if image is not None:
            embeddings["image"] = F.normalize(self.vis_enc(image), p=2, dim=-1)
        if audio is not None:
            embeddings["audio"] = F.normalize(self.aud_enc(audio), p=2, dim=-1)
        return embeddings

# --- Execution & Verification ---
batch_size = 4
model = UnifiedMultimodalSystem(shared_dim=512)

dummy_tokens = torch.randint(0, 1000, (batch_size, 16))      # Text tokens
dummy_images = torch.randn(batch_size, 3, 224, 224)          # RGB Images
dummy_audio  = torch.randn(batch_size, 16000)                # 1-second Audio at 16kHz

outputs = model(text=dummy_tokens, image=dummy_images, audio=dummy_audio)

# Calculate Cross-Modal Similarity between Image 0 and Text 0
img_txt_sim = torch.dot(outputs["image"][0], outputs["text"][0])
print(f"Image-to-Text Cosine Similarity: {img_txt_sim.item():.4f}")
print(f"Shared Representation Shape:    {outputs['image'].shape}")