"""
MultiSense-DF — Lip-Sync Consistency Module
SyncNet-derived dual-stream network: mouth crops + mel-spectrogram
"""

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


class MouthVisualNet(nn.Module):
    """5-layer CNN for 96×96 mouth-region crops."""

    def __init__(self, out_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.ReLU(),   # 48
            nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(),  # 24
            nn.Conv2d(64, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(),# 12
            nn.Conv2d(128, 256, 3, 2, 1), nn.BatchNorm2d(256), nn.ReLU(),# 6
            nn.Conv2d(256, 512, 3, 2, 1), nn.BatchNorm2d(512), nn.ReLU(),# 3
            nn.AdaptiveAvgPool2d(1)
        )
        self.fc = nn.Linear(512, out_dim)

    def forward(self, x):
        # x: (B*T, 3, 96, 96)
        return F.normalize(self.fc(self.net(x).squeeze(-1).squeeze(-1)), dim=-1)


class AudioSpecNet(nn.Module):
    """5-layer CNN for mel-spectrogram windows aligned to mouth frames."""

    def __init__(self, out_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, 256, 3, 2, 1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.Conv2d(256, 512, 3, (2, 1), 1), nn.BatchNorm2d(512), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        self.fc = nn.Linear(512, out_dim)

    def forward(self, x):
        # x: (B*T, 1, 80, W) mel-spectrogram window
        return F.normalize(self.fc(self.net(x).squeeze(-1).squeeze(-1)), dim=-1)


class LipSyncBranch(nn.Module):
    """
    SyncNet-derived lip-sync consistency module.

    Computes per-frame cosine similarity between audio and visual embeddings,
    then summarises temporal synchrony profile with a 2-layer LSTM.

    Input : mouth_crops  (B, T, 3, 96, 96)
            mel_specs    (B, T, 1, 80, W)
    Output: (B, 512) lip-sync embedding + (B, 1) logit
    """

    def __init__(self, sync_dim=256, embed_dim=512, dropout=0.1):
        super().__init__()
        self.visual_net = MouthVisualNet(out_dim=sync_dim)
        self.audio_net  = AudioSpecNet(out_dim=sync_dim)

        # Summarise temporal synchrony profile
        self.lstm = nn.LSTM(
            input_size=sync_dim * 2 + 1,   # [v_feat, a_feat, cos_sim]
            hidden_size=embed_dim // 2,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
            dropout=dropout
        )
        self.norm = nn.LayerNorm(embed_dim)

        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 1)
        )

    def forward(self, mouth_crops, mel_specs):
        B, T, C, H, W = mouth_crops.shape
        _, _, _, F_mel, W_mel = mel_specs.shape

        # Per-frame embeddings
        v = self.visual_net(mouth_crops.reshape(B * T, C, H, W))           # (B*T, 256)
        a = self.audio_net(mel_specs.reshape(B * T, 1, F_mel, W_mel))      # (B*T, 256)

        v = v.reshape(B, T, -1)     # (B, T, 256)
        a = a.reshape(B, T, -1)     # (B, T, 256)

        # Per-frame cosine similarity (scalar)
        cos_sim = F.cosine_similarity(v, a, dim=-1, eps=1e-8).unsqueeze(-1)  # (B, T, 1)

        # Concatenate features + similarity for LSTM input
        seq = torch.cat([v, a, cos_sim], dim=-1)   # (B, T, 513)

        out, _ = self.lstm(seq)          # (B, T, 512)
        embed  = self.norm(out[:, -1])   # last timestep → (B, 512)
        logit  = self.classifier(embed)  # (B, 1)
        return embed, logit, cos_sim.squeeze(-1)  # also return sync profile

    def contrastive_loss(self, mouth_crops, mel_specs, is_synced):
        """
        Contrastive loss for pre-training: in-sync pairs pull together,
        offset pairs push apart.
        is_synced: (B,) bool tensor
        """
        B, T, C, H, W = mouth_crops.shape
        _, _, _, F_mel, W_mel = mel_specs.shape

        v = self.visual_net(mouth_crops.reshape(B * T, C, H, W)).reshape(B, T, -1).mean(1)
        a = self.audio_net(mel_specs.reshape(B * T, 1, F_mel, W_mel)).reshape(B, T, -1).mean(1)

        sim = F.cosine_similarity(v, a, dim=-1)    # (B,)
        margin = 0.4
        loss_pos = (1 - sim) * is_synced.float()
        loss_neg = torch.clamp(sim - margin, min=0) * (~is_synced).float()
        return (loss_pos + loss_neg).mean()


if __name__ == '__main__':
    model = LipSyncBranch()
    mouths = torch.randn(2, 125, 3, 96, 96)
    mels   = torch.randn(2, 125, 1, 80, 16)
    emb, logit, sync = model(mouths, mels)
    print(f'LipSync embed: {emb.shape}, logit: {logit.shape}, sync profile: {sync.shape}')