import torch
import torch.nn as nn
import torchvision
from collections import OrderedDict

NCDICT = {
    'resnet18': 512,
    'resnet34': 512,
    'resnet50': 2048,
    'resnet101': 2048
}

class FrozenBatchNorm2d(nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters are fixed.
    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
    without which any other models than torchvision.models.resnet[18,34,50,101]
    produce nans.
    """

    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(n))
        self.register_buffer("bias", torch.zeros(n))
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]

        super(FrozenBatchNorm2d, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

    def forward(self, x):
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()
        bias = b - rm * scale
        return x * scale + bias
    
class BackboneBase(nn.Module):
    def __init__(self, backbone, num_features):
        super(BackboneBase, self).__init__()
        for name, parameter in backbone.named_parameters():
            if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)
        
        self.body = nn.Sequential(OrderedDict([('conv1', backbone.conv1),
                                               ('bn1', backbone.bn1),
                                               ('relu', backbone.relu),
                                               ('maxpool', backbone.maxpool),
                                               ('layer1', backbone.layer1),
                                               ('layer2', backbone.layer2),
                                               ('layer3', backbone.layer3),
                                               ('layer4', backbone.layer4)]))
        self.num_features = num_features

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

def Resnet(modelname, pretrained=False, norm_layer=FrozenBatchNorm2d):
    model = getattr(torchvision.models, modelname)(pretrained=pretrained, norm_layer=norm_layer)
    model = BackboneBase(model, NCDICT[modelname])
    return model