import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint_sequential
import torchvision.models as models


def build_backbone(name: str):
    if name == "efficientnet_b0":
        m = models.efficientnet_b0(
            weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1
        )
        feat_dim = m.classifier[1].in_features
        m.classifier = nn.Identity()
        original_features = m.features

        def forward_with_ckpt(x):
            return checkpoint_sequential(
                original_features, segments=4, input=x, use_reentrant=False
            )

        m.forward = lambda x: m.avgpool(forward_with_ckpt(x)).flatten(1)
        return m, feat_dim

    elif name == "resnet50":
        m = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        feat_dim = m.fc.in_features
        m.fc = nn.Identity()
        original_layer4 = m.layer4

        def forward_resnet(x):
            x = m.conv1(x); x = m.bn1(x); x = m.relu(x); x = m.maxpool(x)
            x = m.layer1(x); x = m.layer2(x); x = m.layer3(x)
            x = checkpoint_sequential(
                original_layer4, segments=2, input=x, use_reentrant=False
            )
            x = m.avgpool(x)
            return torch.flatten(x, 1)

        m.forward = forward_resnet
        return m, feat_dim

    else:
        return build_backbone("efficientnet_b0")


class VisualPersonalityModel(nn.Module):
    def __init__(self, cfg: dict):
        super().__init__()
        self.backbone, feat_dim = build_backbone(cfg["backbone"])
        agg_dim = feat_dim * 2
        self.projector = nn.Sequential(
            nn.Linear(agg_dim, cfg["embed_dim"] * 2),
            nn.LayerNorm(cfg["embed_dim"] * 2),
            nn.GELU(),
            nn.Dropout(cfg["dropout"]),
            nn.Linear(cfg["embed_dim"] * 2, cfg["embed_dim"]),
            nn.LayerNorm(cfg["embed_dim"]),
        )
        num_traits = len(cfg["traits"])
        self.regressor = nn.Sequential(
            nn.Dropout(cfg["dropout"] * 0.5),
            nn.Linear(cfg["embed_dim"], cfg["embed_dim"] // 2),
            nn.GELU(),
            nn.Dropout(cfg["dropout"] * 0.5),
            nn.Linear(cfg["embed_dim"] // 2, num_traits),
            nn.Sigmoid(),
        )

    def forward(self, frames, backbone_chunk: int = 128):
        B, T, C, H, W = frames.shape
        flat = frames.view(B * T, C, H, W)
        feats = []
        for i in range(0, flat.size(0), backbone_chunk):
            feats.append(self.backbone(flat[i: i + backbone_chunk]))
        x   = torch.cat(feats, dim=0).view(B, T, -1)
        mu  = x.mean(dim=1)
        std = x.std(dim=1).clamp(min=1e-6)
        agg = torch.cat([mu, std], dim=-1)
        emb  = self.projector(agg)
        pred = self.regressor(emb)
        return pred, emb

    def backbone_parameters(self):
        return list(self.backbone.parameters())

    def head_parameters(self):
        return list(self.projector.parameters()) + list(self.regressor.parameters())