"""
All five model classes used in the project.

Each class is copied verbatim from its training notebook so that the saved
state dicts load without any key mismatch. Constructors default to
pretrained=False because the demo overwrites weights from the HF checkpoints.
"""

import timm
import torch
import torch.nn as nn
from torchvision import models


# ===================================================================
# Baseline 1: Xception (frame-level, average per-frame logits)
# Source notebook: 03_baseline_xception.ipynb
# ===================================================================

class XceptionFrameAverager(nn.Module):
    """Run Xception on each frame independently, average the logits to get a clip score."""

    def __init__(self, pretrained=False):
        super().__init__()
        self.backbone = timm.create_model("legacy_xception", pretrained=pretrained, num_classes=1)

    def forward(self, clip):
        # clip: [B, T, 3, H, W]
        batch_size, number_of_frames, channels, height, width = clip.shape

        flat_frames = clip.view(batch_size * number_of_frames, channels, height, width)
        per_frame_logits = self.backbone(flat_frames)              # [B*T, 1]

        per_frame_logits = per_frame_logits.view(batch_size, number_of_frames, -1)
        clip_logits = per_frame_logits.mean(dim=1)                 # [B, 1]
        return clip_logits


# ===================================================================
# Baseline 2: EfficientNet-B3 (frame-level, average per-frame logits)
# Source notebook: 04_baseline_efficientnet.ipynb
# ===================================================================

class EfficientNetFrameAverager(nn.Module):
    """Run EfficientNet-B3 on each frame independently, average the logits to get a clip score."""

    def __init__(self, pretrained=False):
        super().__init__()
        self.backbone = timm.create_model("tf_efficientnet_b3", pretrained=pretrained, num_classes=1)

    def forward(self, clip):
        batch_size, number_of_frames, channels, height, width = clip.shape

        flat_frames = clip.view(batch_size * number_of_frames, channels, height, width)
        per_frame_logits = self.backbone(flat_frames)              # [B*T, 1]

        per_frame_logits = per_frame_logits.view(batch_size, number_of_frames, -1)
        clip_logits = per_frame_logits.mean(dim=1)                 # [B, 1]
        return clip_logits


# ===================================================================
# Baseline 3: CNN (ResNet-18) + 2-layer BiLSTM
# Source notebook: 05_baseline_cnn_lstm.ipynb
# ===================================================================

class CnnLstmClipModel(nn.Module):
    """ResNet-18 features per frame, then a 2-layer BiLSTM over the 24 frames, then a binary head."""

    def __init__(self, lstm_hidden_size=256, lstm_num_layers=2, lstm_dropout=0.3, pretrained=False):
        super().__init__()
        resnet_weights = models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None
        resnet = models.resnet18(weights=resnet_weights)
        self.feature_dim = resnet.fc.in_features                   # 512 for ResNet-18
        resnet.fc = nn.Identity()
        self.cnn_feature_extractor = resnet

        self.bilstm = nn.LSTM(
            input_size=self.feature_dim,
            hidden_size=lstm_hidden_size,
            num_layers=lstm_num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=lstm_dropout if lstm_num_layers > 1 else 0.0,
        )

        self.classifier = nn.Sequential(
            nn.Linear(lstm_hidden_size * 2, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 1),
        )

    def forward(self, clip):
        batch_size, number_of_frames, channels, height, width = clip.shape

        flat_frames = clip.view(batch_size * number_of_frames, channels, height, width)
        per_frame_features = self.cnn_feature_extractor(flat_frames)
        per_frame_features = per_frame_features.view(batch_size, number_of_frames, self.feature_dim)

        lstm_output, (final_hidden, _) = self.bilstm(per_frame_features)
        # final_hidden layout: [layer0_fwd, layer0_bwd, layer1_fwd, layer1_bwd, ...]
        forward_last = final_hidden[-2]
        backward_last = final_hidden[-1]
        clip_representation = torch.cat([forward_last, backward_last], dim=1)

        clip_logits = self.classifier(clip_representation)         # [B, 1]
        return clip_logits


# ===================================================================
# Baseline 4: ViT-Base/16 (frame-level, average per-frame logits)
# Source notebook: 06_baseline_vit.ipynb
# ===================================================================

class ViTFrameAverager(nn.Module):
    """Run ViT-Base/16 on each frame independently, average the logits to get a clip score."""

    def __init__(self, pretrained=False):
        super().__init__()
        self.backbone = timm.create_model("vit_base_patch16_224", pretrained=pretrained, num_classes=1)

    def forward(self, clip):
        batch_size, number_of_frames, channels, height, width = clip.shape

        flat_frames = clip.view(batch_size * number_of_frames, channels, height, width)
        per_frame_logits = self.backbone(flat_frames)              # [B*T, 1]

        per_frame_logits = per_frame_logits.view(batch_size, number_of_frames, -1)
        clip_logits = per_frame_logits.mean(dim=1)                 # [B, 1]
        return clip_logits


# ===================================================================
# Proposed: Hybrid CNN-Transformer
# Source notebook: 07_proposed_hybrid_cnn_transformer.ipynb
# ===================================================================

class HybridCNNTransformer(nn.Module):
    """EfficientNet-B3 per-frame backbone, then a Transformer encoder over the
    frame sequence with a CLS token."""

    def __init__(
        self,
        pretrained=False,
        number_of_frames=24,
        number_of_transformer_layers=4,
        number_of_heads=8,
        transformer_feedforward_dimensions=2048,
        dropout=0.1,
    ):
        super().__init__()

        # num_classes=0 makes timm return the pre-classifier pooled feature
        # (1536-dim for EfficientNet-B3).
        self.backbone = timm.create_model(
            "tf_efficientnet_b3", pretrained=pretrained, num_classes=0
        )
        feature_dimensions = self.backbone.num_features            # 1536

        self.cls_token = nn.Parameter(torch.zeros(1, 1, feature_dimensions))
        nn.init.trunc_normal_(self.cls_token, std=0.02)

        self.positional_embeddings = nn.Parameter(
            torch.zeros(1, number_of_frames + 1, feature_dimensions)
        )
        nn.init.trunc_normal_(self.positional_embeddings, std=0.02)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=feature_dimensions,
            nhead=number_of_heads,
            dim_feedforward=transformer_feedforward_dimensions,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=number_of_transformer_layers
        )
        self.transformer_norm = nn.LayerNorm(feature_dimensions)

        self.classifier = nn.Sequential(
            nn.Linear(feature_dimensions, feature_dimensions // 2),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(feature_dimensions // 2, 1),
        )

    def forward(self, clip):
        batch_size, number_of_frames, channels, height, width = clip.shape

        flat_frames = clip.view(batch_size * number_of_frames, channels, height, width)
        frame_features = self.backbone(flat_frames)
        frame_features = frame_features.view(batch_size, number_of_frames, -1)

        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        sequence = torch.cat([cls_tokens, frame_features], dim=1)
        sequence = sequence + self.positional_embeddings

        sequence = self.transformer_encoder(sequence)
        sequence = self.transformer_norm(sequence)
        cls_output = sequence[:, 0]
        clip_logits = self.classifier(cls_output)
        return clip_logits