# Module 6: End-to-End PyTorch Implementation
# Below is a complete, modular, and pedagogical implementation of a tri-modal embedding fusion and metric alignment pipeline in PyTorch.
# Read through the classes to trace how shapes propagate through linear projections and attention heads.

import torch
import torch.nn as nn
import torch.nn.functional as F
class TextEncoder(nn.Module):
    def __init__(self, vocab_size=5000, hidden_dim=256, embed_dim=128):
        super().__init__()
        # Converts integer tokens to vectors
        self.token_embeddings = nn.Embedding(vocab_size, hidden_dim)
        # Simplified Transformer-like processing block
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim, nhead=4, dim_feedforward=512, batch_first=True
        )
        # Projections to share joint space
        self.projection_head = nn.Linear(hidden_dim, embed_dim)

    def forward(self, x):
        # x shape: [Batch, SeqLength]
        x_emb = self.token_embeddings(x) # [Batch, SeqLength, HiddenDim]
        x_trans = self.encoder_layer(x_emb) # [Batch, SeqLength, HiddenDim]
        # Pool sequence tokens via mean-pooling across temporal sequence
        x_pooled = x_trans.mean(dim=1) # [Batch, HiddenDim]
        # Project to joint metric dimension & L2-normalize
        z = self.projection_head(x_pooled) # [Batch, EmbedDim]
        return F.normalize(z, p=2, dim=-1)

class VisionEncoder(nn.Module):
    def __init__(self, patch_dim=16, in_channels=3, hidden_dim=256, embed_dim=128):
        super().__init__()
        # Convolutional Patch Extraction simulating a ViT projection
        self.patch_conv = nn.Conv2d(
            in_channels, hidden_dim, kernel_size=patch_dim, stride=patch_dim
        )
        self.projection_head = nn.Linear(hidden_dim, embed_dim)

    def forward(self, x):
        # x shape: [Batch, Channels, Height, Width]
        patches = self.patch_conv(x) # [Batch, HiddenDim, PatchesH, PatchesW]
        patches = patches.flatten(2).transpose(1, 2) # [Batch, N_Patches, HiddenDim]
        # Pool spatially via mean across patches
        img_pooled = patches.mean(dim=1) # [Batch, HiddenDim]
        # Project & L2-normalize
        z = self.projection_head(img_pooled) # [Batch, EmbedDim]
        return F.normalize(z, p=2, dim=-1)

class AudioEncoder(nn.Module):
    def __init__(self, hidden_dim=256, embed_dim=128):
        super().__init__()
        # 2D Convolutional extraction from Log-Mel Spectrogram frames
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.fc = nn.Linear(64, hidden_dim)
        self.projection_head = nn.Linear(hidden_dim, embed_dim)

    def forward(self, spec):
        # spec shape: [Batch, 1, MelBins, TimeSteps]
        feat = F.relu(self.conv1(spec))
        feat = F.relu(self.conv2(feat))
        # Global average pooling across remaining 2D spatial dimensions
        feat_pooled = feat.mean(dim=[2, 3]) # [Batch, 64]
        feat_hidden = F.relu(self.fc(feat_pooled)) # [Batch, HiddenDim]
        # Project & L2-normalize
        z = self.projection_head(feat_hidden) # [Batch, EmbedDim]
        return F.normalize(z, p=2, dim=-1)

class UnifiedMultimodalSystem(nn.Module):
    def __init__(self, embed_dim=128):
        super().__init__()
        self.text_encoder = TextEncoder(embed_dim=embed_dim)
        self.vision_encoder = VisionEncoder(embed_dim=embed_dim)
        self.audio_encoder = AudioEncoder(embed_dim=embed_dim)
        self.temp = nn.Parameter(torch.ones([]) * 0.07) # temperature parameter

    def compute_infonce_loss(self, z_a, z_b):
        # Symmetric InfoNCE Calculation
        batch_size = z_a.size(0)
        # Calculate cosine similarity matrix explicitly via scaled dot-product
        sim_matrix = torch.matmul(z_a, z_b.t()) / self.temp
        labels = torch.arange(batch_size, device=z_a.device)

        loss_a = F.cross_entropy(sim_matrix, labels)
        loss_b = F.cross_entropy(sim_matrix.t(), labels)
        return (loss_a + loss_b) / 2.0

    def forward(self, text, image, audio):
        z_t = self.text_encoder(text)
        z_v = self.vision_encoder(image)
        z_a = self.audio_encoder(audio)

        # Multi-way alignment losses
        loss_tv = self.compute_infonce_loss(z_t, z_v)
        loss_ta = self.compute_infonce_loss(z_t, z_a)
        loss_va = self.compute_infonce_loss(z_v, z_a)

        total_loss = (loss_tv + loss_ta + loss_va) / 3.0
        return total_loss, z_t, z_v, z_a

# Dry-Run Validation Code Script
# Run the verification block below to check dimensions, trace shape transitions, and inspect output tensors of the system on a simulated mini-batch:

if __name__ == '__main__':
    model = UnifiedMultimodalSystem(embed_dim=128)

    # Simulated Mini-Batch of Size N=2
    mock_text = torch.randint(0, 5000, (2, 10)) # [Batch=2, SeqLength=10]
    mock_image = torch.randn(2, 3, 224, 224)    # [Batch=2, C=3, H=224, W=224]
    mock_audio = torch.randn(2, 1, 64, 100)     # [Batch=2, C=1, Mel=64, Time=100]

    loss, z_t, z_v, z_a = model(mock_text, mock_image, mock_audio)

    print('=== PIPELINE DIMENSION SHAPES ===')
    print(f'Text Latent Shape:   {z_t.shape} (Expected: [2, 128])')
    print(f'Vision Latent Shape: {z_v.shape} (Expected: [2, 128])')
    print(f'Audio Latent Shape:  {z_a.shape} (Expected: [2, 128])')
    print(f'Unified Joint Loss:  {loss.item():.4f}')