import torch
import torch.nn as nn
from torchvision.models.resnet import BasicBlock, Bottleneck, conv1x1, conv3x3
from torchvision.models.utils import load_state_dict_from_url
import torchvision

from models.arch_util import ConvBnRelu
from models.pixel_level_contrastive_learning.resnet_unet import ReverseBottleneck
from trainer.networks import register_model
from utils.util import checkpoint, opt_get


class UResNet50_3(torchvision.models.resnet.ResNet):

    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
                 norm_layer=None, out_dim=128):
        super().__init__(block, layers, num_classes, zero_init_residual, groups, width_per_group,
                         replace_stride_with_dilation, norm_layer)
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        '''
        # For reference:
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
                                       dilate=replace_stride_with_dilation[2])
        '''
        uplayers = []
        inplanes = 2048
        first = True
        for i in range(3):
            uplayers.append(ReverseBottleneck(inplanes, inplanes // 2, norm_layer=norm_layer, passthrough=not first))
            inplanes = inplanes // 2
            first = False
        self.uplayers = nn.ModuleList(uplayers)

        # These two variables are separated out and renamed so that I can re-use parameters from a pretrained resnet_unet2.
        self.last_uplayer = ReverseBottleneck(256, 128, norm_layer=norm_layer, passthrough=True)
        self.tail3 = nn.Sequential(conv1x1(192, 128),
                                  norm_layer(128),
                                  nn.ReLU(),
                                  conv1x1(128, out_dim))

        del self.fc  # Not used in this implementation and just consumes a ton of GPU memory.


    def _forward_impl(self, x):
        x0 = self.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x0)

        x1 = checkpoint(self.layer1, x)
        x2 = checkpoint(self.layer2, x1)
        x3 = checkpoint(self.layer3, x2)
        x4 = checkpoint(self.layer4, x3)
        unused = self.avgpool(x4)  # This is performed for instance-level pixpro learning, even though it is unused.

        x = checkpoint(self.uplayers[0], x4)
        x = checkpoint(self.uplayers[1], x, x3)
        x = checkpoint(self.uplayers[2], x, x2)
        x = checkpoint(self.last_uplayer, x, x1)

        return checkpoint(self.tail3, torch.cat([x, x0], dim=1))

    def forward(self, x):
        return self._forward_impl(x)


@register_model
def register_u_resnet50_3(opt_net, opt):
    model = UResNet50_3(Bottleneck, [3, 4, 6, 3], out_dim=opt_net['odim'])
    if opt_get(opt_net, ['use_pretrained_base'], False):
        state_dict = load_state_dict_from_url('https://download.pytorch.org/models/resnet50-19c8e357.pth', progress=True)
        model.load_state_dict(state_dict, strict=False)
    return model


if __name__ == '__main__':
    model = UResNet50_3(Bottleneck, [3,4,6,3])
    samp = torch.rand(1,3,224,224)
    y = model(samp)
    print(y.shape)
    # For pixpro: attach to "tail.3"