import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

from rscd.models.backbones.Decom_Backbone import ResNet3D


class AFCD3D_backbone(nn.Module):
    def __init__(self):
        super(AFCD3D_backbone, self).__init__()
        resnet = torchvision.models.resnet18(pretrained=True)
        self.resnet = ResNet3D(resnet)

    def forward(self, imageA, imageB):
        imageA = imageA.unsqueeze(2)
        imageB = imageB.unsqueeze(2)
        x = torch.cat([imageA, imageB], 2)
        size = x.size()[3:]
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x0 = self.resnet.relu(x)
        x = self.resnet.maxpool(x0)
        x1 = self.resnet.layer1(x)
        x2 = self.resnet.layer2(x1)
        x3 = self.resnet.layer3(x2)
        x4 = self.resnet.layer4(x3)

        return [size, x0, x1, x2, x3, x4]