import torch
import torch.nn as nn
from src.anoseg.models.DFC.backbone.vgg import vgg19

"""
    Code based on https://github.com/YoungGod/DFC
    Paper: https://doi.org/10.1016/j.patcog.2022.108874
"""


# Res blocks
class ResnetBlock(nn.Module):
    def __init__(self, dim, dilation=2, padding='reflect'):
        super(ResnetBlock, self).__init__()
        self.conv_block = nn.Sequential(
            nn.ReflectionPad2d(dilation) if padding == 'reflect' else nn.ZeroPad2d(dilation),
            nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=dilation),
            # models.InstanceNorm2d(dim, track_running_stats=False),
            nn.ReLU(True),

            nn.ReflectionPad2d(1) if padding == 'reflect' else nn.ZeroPad2d(1),
            nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=1),
            # models.InstanceNorm2d(dim, track_running_stats=False),
        )

    def forward(self, x):
        out = x + self.conv_block(x)  # ReLU or not?
        # print(out.size())   ##########################################
        # Remove ReLU at the end of the residual block
        # http://torch.ch/blog/2016/02/04/resnets.html

        return out


# current Estimator is using 4 residual blocks
class MCRL(nn.Module):
    def __init__(self, in_channels=256, residual_blocks=2, dilation=2):
        super(MCRL, self).__init__()

        # deformable branch (data dependent context)
        blocks = []
        for _ in range(residual_blocks):
            block = ResnetBlock(in_channels, dilation=2, padding='reflect')
            # block = ResnetBlock_BDCM(in_channels, dilation=dilation)
            blocks.append(block)
        self.middle_res1 = nn.Sequential(*blocks)

        # dilated branch (fixed context)
        blocks = []
        for _ in range(residual_blocks):
            block = ResnetBlock(in_channels, dilation=4, padding='reflect')
            blocks.append(block)
        self.middle_res2 = nn.Sequential(*blocks)

        # dilated branch (fixed context)
        blocks = []
        for _ in range(residual_blocks):
            block = ResnetBlock(in_channels, dilation=8, padding='reflect')
            blocks.append(block)
        self.middle_res3 = nn.Sequential(*blocks)

    def forward(self, x):
        x1 = self.middle_res1(x)
        x2 = self.middle_res2(x)
        x3 = self.middle_res3(x)
        out = x + x1 + x2 + x3
        return out


class VGG19_S(torch.nn.Module):
    """
    VGG19: [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
    level1: 64*2=128; level2: 128*2=256; level3: 256*4=1024; level4: 512*4=2048; level5: 512*4=2048
    Total dimension: 128 + 256 + 1024 + 2048 + 2048 = 5504
    """

    def __init__(self, pretrain=False, gradient=False, pool='avg'):
        super(VGG19_S, self).__init__()
        features = vgg19(pretrained=pretrain, progress=False).features  # feature layers
        """ vgg.features
        Sequential(
          (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (1): ReLU(inplace)                                                        # self.relu1_1
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (3): ReLU(inplace)                                                        # self.relu1_2
          
          (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (6): ReLU(inplace)
          (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (8): ReLU(inplace)  
              
          (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (11): ReLU(inplace)
          (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (13): ReLU(inplace)
          (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (15): ReLU(inplace)
          (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (17): ReLU(inplace)
          
          (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (20): ReLU(inplace)
          (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (22): ReLU(inplace)
          (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (24): ReLU(inplace)
          (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (26): ReLU(inplace)
          
          (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (29): ReLU(inplace)
          (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (31): ReLU(inplace)
          (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (33): ReLU(inplace)
          (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (35): ReLU(inplace)
          
          (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        )
        """
        # hierarchy 1 (level 1)
        self.conv1_1 = features[0]
        self.relu1_1 = features[1]
        self.conv1_2 = features[2]
        self.relu1_2 = features[3]

        # hierarchy 2 (level 2)
        self.pool1 = features[4]
        self.conv2_1 = features[5]
        self.relu2_1 = features[6]
        self.conv2_2 = features[7]
        self.relu2_2 = features[8]

        # replace all convs in level 3 with deformable conv v2
        # hierarchy 3 (level 3)    
        self.pool2 = features[9]
        self.conv3_1 = features[10]
        self.relu3_1 = features[11]
        self.conv3_2 = features[12]
        self.relu3_2 = features[13]
        self.conv3_3 = features[14]
        self.relu3_3 = features[15]
        self.conv3_4 = features[16]
        self.relu3_4 = features[17]

        # hierarchy 4 (level 4)
        self.pool3 = features[18]
        self.conv4_1 = features[19]
        self.relu4_1 = features[20]
        self.conv4_2 = features[21]
        self.relu4_2 = features[22]
        self.conv4_3 = features[23]
        self.relu4_3 = features[24]
        self.conv4_4 = features[25]
        self.relu4_4 = features[26]

        # after pool3
        self.resnet = MCRL(in_channels=self.conv3_4.in_channels, residual_blocks=4, dilation=2)

        # self.resnet = DilatedDenseResidualBlock(in_channels=self.conv3_4.in_channels, dilated_dense_blocks=2, growthrate=32, num_layer=4, dilation=True)
        # # init as zero
        # def weight_init(m):
        #     """
        #     ref: https://blog.csdn.net/qq_36338754/article/details/97756378
        #     """
        #     if isinstance(m, models.Conv2d) or isinstance(m, models.Linear):
        #         if m.weight is not None:
        #             m.weight.data.zero_()
        #         if m.bias is not None:
        #             m.bias.data.zero_()
        # self.resnet.apply(weight_init)    

        # hierarchy 5 (level 5)
        self.pool4 = features[27]
        self.conv5_1 = features[28]
        self.relu5_1 = features[29]
        self.conv5_2 = features[30]
        self.relu5_2 = features[31]
        self.conv5_3 = features[32]
        self.relu5_3 = features[33]
        self.conv5_4 = features[34]
        self.relu5_4 = features[35]

        self.pool5 = features[36]

        # don't need the gradients, just want the features
        if not gradient:
            for param in self.parameters():
                param.requires_grad = False

        # reflection padding
        self.pad = nn.ReflectionPad2d(padding=1)

        # avg pooling
        if pool == 'avg':
            avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
            # self.pool3 = self.pool4 = self.pool5 = avg_pool
            self.pool1 = self.pool2 = self.pool3 = self.pool4 = self.pool5 = avg_pool

    def forward(self, x, feature_layers):
        # level 1    pad before conv
        x = self.pad(x)
        conv1_1 = self.conv1_1(x)
        relu1_1 = self.relu1_1(conv1_1)
        conv1_2 = self.conv1_2(self.pad(relu1_1))
        relu1_2 = self.relu1_2(conv1_2)
        pool1 = self.pool1(relu1_2)

        # level 2
        pool1 = self.pad(pool1)
        conv2_1 = self.conv2_1(pool1)
        relu2_1 = self.relu2_1(conv2_1)
        conv2_2 = self.conv2_2(self.pad(relu2_1))
        relu2_2 = self.relu2_2(conv2_2)
        pool2 = self.pool2(relu2_2)

        # level 3
        pool2 = self.pad(pool2)
        conv3_1 = self.conv3_1(pool2)
        relu3_1 = self.relu3_1(conv3_1)
        conv3_2 = self.conv3_2(self.pad(relu3_1))
        relu3_2 = self.relu3_2(conv3_2)
        conv3_3 = self.conv3_3(self.pad(relu3_2))
        # conv3_3 = self.conv3_3(relu3_2)
        relu3_3 = self.relu3_3(conv3_3)
        conv3_4 = self.conv3_4(self.pad(relu3_3))
        relu3_4 = self.relu3_4(conv3_4)
        pool3 = self.pool3(relu3_4)

        # level 4
        pool3 = self.resnet(pool3)  # resblock
        pool3 = self.pad(pool3)
        conv4_1 = self.conv4_1(pool3)
        relu4_1 = self.relu4_1(conv4_1)
        conv4_2 = self.conv4_2(self.pad(relu4_1))
        relu4_2 = self.relu4_2(conv4_2)
        conv4_3 = self.conv4_3(self.pad(relu4_2))
        relu4_3 = self.relu4_3(conv4_3)
        conv4_4 = self.conv4_4(self.pad(relu4_3))
        relu4_4 = self.relu4_4(conv4_4)
        pool4 = self.pool4(relu4_4)

        # level 5
        pool4 = self.pad(pool4)
        conv5_1 = self.conv5_1(pool4)
        relu5_1 = self.relu5_1(conv5_1)
        conv5_2 = self.conv5_2(self.pad(relu5_1))
        relu5_2 = self.relu5_2(conv5_2)
        # conv5_3 = self.conv5_3(self.pad(relu5_2))
        # relu5_3 = self.relu5_3(conv5_3)
        # conv5_4 = self.conv5_4(self.pad(relu5_3))
        # relu5_4 = self.relu5_4(conv5_4)
        # pool5 = self.pool5(relu5_4)

        out = {
            'relu1_1': relu1_1,
            'relu1_2': relu1_2,

            'relu2_1': relu2_1,
            'relu2_2': relu2_2,

            'relu3_1': relu3_1,
            'relu3_2': relu3_2,
            'relu3_3': relu3_3,
            'relu3_4': relu3_4,

            'relu4_1': relu4_1,
            'relu4_2': relu4_2,
            'relu4_3': relu4_3,
            'relu4_4': relu4_4,

            'relu5_1': relu5_1,
            'relu5_2': relu5_2,
            # 'relu5_3': relu5_3,
            # 'relu5_4': relu5_4,
        }
        return dict((key, value) for key, value in out.items() if key in feature_layers)


if __name__ == '__main__':
    x = torch.randn((1, 3, 256, 256), dtype=torch.float32)
    model = VGG19_S()
    y = model(x, feature_layers=("relu1_1", "relu1_2", "relu2_1", "relu2_2",
                                 "relu3_1", "relu3_2", "relu3_3", "relu3_4",
                                 "relu4_1", "relu4_2", "relu4_3", "relu4_4",
                                 "relu5_1", "relu5_2"))  # , "relu5_3", "relu5_4"
    for key, value in y.items():
        print(key, value.shape)
    print(model.conv3_1.in_channels, model.conv3_4.out_channels)
    print(model.conv4_1.in_channels, model.conv4_2.out_channels)
    print(model.conv4_1.kernel_size, model.conv4_2.out_channels)
