"""
MultiSense-DF — Visual Branch
EfficientNet-B4 CNN backbone + Temporal Transformer
"""

import torch
import torch.nn as nn
from einops import rearrange
import timm


class TemporalTransformer(nn.Module):
    """6-layer Transformer over frame-level embeddings."""

    def __init__(self, embed_dim=512, num_heads=8, num_layers=6,
                 num_frames=125, dropout=0.1):
        super().__init__()
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_frames + 1, embed_dim))

        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.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.norm = nn.LayerNorm(embed_dim)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)

    def forward(self, x):
        # x: (B, T, D)
        B = x.size(0)
        cls = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls, x], dim=1)          # (B, T+1, D)
        x = x + self.pos_embed[:, :x.size(1)]
        x = self.transformer(x)
        x = self.norm(x)
        return x[:, 0]                            # CLS token → (B, D)


class VisualBranch(nn.Module):
    """
    EfficientNet-B4 frame encoder + TemporalTransformer.
    Input : (B, T, 3, 224, 224) video clips
    Output: (B, 512) visual embedding + (B, 1) logit
    """

    def __init__(self, embed_dim=512, num_frames=125,
                 num_heads=8, num_layers=6, dropout=0.1,
                 freeze_backbone_epochs=5):
        super().__init__()
        self.freeze_backbone_epochs = freeze_backbone_epochs

        # EfficientNet-B4 backbone (1792-d features)
        self.backbone = timm.create_model(
            'efficientnet_b4', pretrained=True, num_classes=0
        )
        backbone_dim = self.backbone.num_features   # 1792

        # Project backbone → embed_dim
        self.proj = nn.Sequential(
            nn.Linear(backbone_dim, embed_dim),
            nn.LayerNorm(embed_dim),
            nn.GELU()
        )

        # Temporal Transformer
        self.temporal = TemporalTransformer(
            embed_dim=embed_dim,
            num_heads=num_heads,
            num_layers=num_layers,
            num_frames=num_frames,
            dropout=dropout
        )

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

    def freeze_backbone(self):
        for p in self.backbone.parameters():
            p.requires_grad = False

    def unfreeze_backbone(self):
        for p in self.backbone.parameters():
            p.requires_grad = True

    def forward(self, x):
        # x: (B, T, C, H, W)
        B, T, C, H, W = x.shape
        frames = rearrange(x, 'b t c h w -> (b t) c h w')

        # Extract per-frame features
        feats = self.backbone(frames)              # (B*T, 1792)
        feats = self.proj(feats)                   # (B*T, 512)
        feats = rearrange(feats, '(b t) d -> b t d', b=B, t=T)

        # Temporal aggregation
        embed = self.temporal(feats)               # (B, 512)
        logit = self.classifier(embed)             # (B, 1)
        return embed, logit


if __name__ == '__main__':
    model = VisualBranch(num_frames=125)
    dummy = torch.randn(2, 125, 3, 224, 224)
    emb, logit = model(dummy)
    print(f'Visual embed: {emb.shape}, logit: {logit.shape}')