"""
MultiSense-DF — Cross-Modal Attention Fusion Module
Fuses visual, audio, and lip-sync embeddings via Transformer self-attention
"""

import torch
import torch.nn as nn


class CrossModalFusion(nn.Module):
    """
    Takes 3 modality embeddings (visual, audio, lip-sync), each 512-d,
    stacks them as a 3-token sequence, and applies 2 layers of multi-head
    self-attention so each modality can attend to the others.

    Per-modality auxiliary heads + global classification head.

    Input : vis_emb   (B, 512)
            aud_emb   (B, 512)
            sync_emb  (B, 512)
    Output:
        global_logit  (B, 1)
        per_mod_logits: dict {'visual': (B,1), 'audio': (B,1), 'lipsync': (B,1)}
        attn_weights  list of attention weight tensors per layer
    """

    def __init__(self, embed_dim=512, num_heads=8, num_fusion_layers=2,
                 num_modalities=3, dropout=0.1):
        super().__init__()
        self.num_modalities = num_modalities

        # Learnable modality-type embeddings (like segment embeddings in BERT)
        self.modality_embed = nn.Parameter(
            torch.zeros(1, num_modalities, embed_dim)
        )
        nn.init.trunc_normal_(self.modality_embed, std=0.02)

        # Fusion Transformer (2 layers, 8 heads)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=num_heads,
            dim_feedforward=embed_dim * 4,
            dropout=dropout, batch_first=True, norm_first=True
        )
        self.fusion_transformer = nn.TransformerEncoder(
            encoder_layer, num_layers=num_fusion_layers,
            enable_nested_tensor=False
        )
        self.norm = nn.LayerNorm(embed_dim)

        # Per-modality classification heads (auxiliary)
        self.vis_head  = self._make_head(embed_dim, dropout)
        self.aud_head  = self._make_head(embed_dim, dropout)
        self.sync_head = self._make_head(embed_dim, dropout)

        # Global classification head (mean-pool over modality tokens)
        self.global_head = self._make_head(embed_dim, dropout)

    @staticmethod
    def _make_head(embed_dim, dropout):
        return nn.Sequential(
            nn.Linear(embed_dim, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 1)
        )

    def forward(self, vis_emb, aud_emb, sync_emb):
        # Stack: (B, 3, 512)
        tokens = torch.stack([vis_emb, aud_emb, sync_emb], dim=1)
        tokens = tokens + self.modality_embed   # add modality type embedding

        # Cross-modal self-attention
        fused = self.fusion_transformer(tokens) # (B, 3, 512)
        fused = self.norm(fused)

        vis_out  = fused[:, 0]   # (B, 512)
        aud_out  = fused[:, 1]
        sync_out = fused[:, 2]
        global_  = fused.mean(1) # mean-pool

        global_logit = self.global_head(global_)

        per_mod_logits = {
            'visual':  self.vis_head(vis_out),
            'audio':   self.aud_head(aud_out),
            'lipsync': self.sync_head(sync_out)
        }

        # Attention weights for explainability (hook-based — approximated here)
        attn_weights = {
            'visual_weight':  vis_out.norm(dim=-1, keepdim=True),
            'audio_weight':   aud_out.norm(dim=-1, keepdim=True),
            'lipsync_weight': sync_out.norm(dim=-1, keepdim=True),
        }

        return global_logit, per_mod_logits, attn_weights


class MultiSenseDF(nn.Module):
    """
    Complete MultiSense-DF system.
    Composes all three branches + cross-modal fusion.
    """

    def __init__(self, visual_branch, audio_branch, lipsync_branch,
                 embed_dim=512, num_heads=8, dropout=0.1):
        super().__init__()
        self.visual  = visual_branch
        self.audio   = audio_branch
        self.lipsync = lipsync_branch
        self.fusion  = CrossModalFusion(
            embed_dim=embed_dim, num_heads=num_heads,
            dropout=dropout
        )

    def forward(self, frames, waveform, mouth_crops, mel_specs,
                audio_mask=None):
        vis_emb,  vis_logit  = self.visual(frames)
        aud_emb,  aud_logit  = self.audio(waveform, audio_mask)
        sync_emb, sync_logit, sync_profile = self.lipsync(mouth_crops, mel_specs)

        global_logit, per_mod_logits, attn_weights = self.fusion(
            vis_emb, aud_emb, sync_emb
        )

        return {
            'global_logit':   global_logit,
            'per_mod_logits': per_mod_logits,
            'branch_logits': {
                'visual':  vis_logit,
                'audio':   aud_logit,
                'lipsync': sync_logit
            },
            'sync_profile': sync_profile,
            'attn_weights': attn_weights
        }


if __name__ == '__main__':
    from visual_branch  import VisualBranch
    from audio_branch   import AudioBranch
    from lipsync_branch import LipSyncBranch

    vis   = VisualBranch(num_frames=125)
    aud   = AudioBranch()
    sync  = LipSyncBranch()
    model = MultiSenseDF(vis, aud, sync)

    frames     = torch.randn(1, 125, 3, 224, 224)
    waveform   = torch.randn(1, 80000)
    mouths     = torch.randn(1, 125, 3, 96, 96)
    mels       = torch.randn(1, 125, 1, 80, 16)

    out = model(frames, waveform, mouths, mels)
    print('Global logit:', out['global_logit'].shape)
    print('Per-mod logits:', {k: v.shape for k, v in out['per_mod_logits'].items()})