import torch
import torch.nn as nn
import timm


class DeepfakeModel(nn.Module):

    def __init__(self):
        super().__init__()

        self.cnn = timm.create_model(
            "tf_efficientnetv2_b1",
            pretrained=False,
            num_classes=0,
            global_pool="avg",
            drop_rate=0.2
        )

        feat_dim = self.cnn.num_features

        self.proj = nn.Linear(
            feat_dim,
            128
        )

        encoder = nn.TransformerEncoderLayer(
            d_model=128,
            nhead=2,
            batch_first=True,
            dropout=0.2
        )

        self.transformer = nn.TransformerEncoder(
            encoder,
            num_layers=1
        )

        self.dropout = nn.Dropout(
            0.25
        )

        self.fc = nn.Linear(
            128,
            2
        )

    def forward(self, x):

        B, T, C, H, W = x.shape

        x = x.view(
            B * T,
            C,
            H,
            W
        )

        x = self.cnn(x)

        x = x.view(
            B,
            T,
            -1
        )

        x = self.proj(x)

        x = self.transformer(x)

        x = x.mean(dim=1)

        x = self.dropout(x)

        return self.fc(x)