import torch
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin
from torch import nn
from torchvision import models


class ICN(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super().__init__()

        cnn = models.resnet50(pretrained=False)
        self.cnn_head = nn.Sequential(
            *list(cnn.children())[:4],
            *list(list(list(cnn.children())[4].children())[0].children())[:4],
        )
        self.cnn_tail = nn.Sequential(
            *list(list(cnn.children())[4].children()
                  )[1:], *list(cnn.children())[5:-2]
        )

        self.conv1 = nn.Conv2d(128, 256, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(num_features=256)

        self.fc1 = nn.Linear(2048 * 7 * 7, 256)
        self.fc2 = nn.Linear(256, 7 * 7)

        self.cls_fc = nn.Linear(256, 3)

        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        # Input: [-1, 6, 224, 224]
        real = x[:, :3, :, :]
        fake = x[:, 3:, :, :]

        # Push both images through pretrained backbone
        real_features = F.relu(self.cnn_head(real))  # [-1, 64, 56, 56]
        fake_features = F.relu(self.cnn_head(fake))  # [-1, 64, 56, 56]

        # [-1, 128, 56, 56]
        combined = torch.cat((real_features, fake_features), 1)

        x = self.conv1(combined)  # [-1, 256, 56, 56]
        x = self.bn1(x)
        x = F.relu(x)

        x = self.cnn_tail(x)
        x = x.view(-1, 2048 * 7 * 7)

        # Final feature [-1, 256]
        d = F.relu(self.fc1(x))

        # Heatmap [-1, 49]
        grid = self.fc2(d)

        # Classifier [-1, 1]
        cl = self.cls_fc(d)

        return grid, cl