# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Backbone modules.
"""
from collections import OrderedDict

import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from typing import Dict, List

from util.misc import NestedTensor, is_main_process


class FrozenBatchNorm2d(torch.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: nn.Module, train_backbone: bool, return_interm_layers: bool):
        super().__init__()
        # for name, parameter in backbone.named_parameters():
        #     if not train_backbone or 'layer1' not in name and 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
        #         parameter.requires_grad_(False)
        if return_interm_layers:
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
            self.num_channels = [256, 512, 1024, 2048]
            # self.num_channels = [128, 256, 512]  # for resnet18
        else:
            return_layers = {'layer3': "0"}
            self.num_channels = [1024]
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)

    def forward(self, tensor_list):
        xs = self.body(tensor_list)
        out = [xs[str(i)] for i in range(len(xs))]
        return out


class Backbone(BackboneBase):
    """ResNet backbone with frozen BatchNorm."""
    def __init__(self, name: str,
                 train_backbone: bool,
                 return_interm_layers: bool):
        backbone = getattr(torchvision.models, name)(
            replace_stride_with_dilation=[False, False, False],  # NOTICE!
            pretrained=is_main_process(), norm_layer=None)
        backbone.maxpool = nn.Identity()
        super().__init__(backbone, train_backbone, return_interm_layers)


class HourglassBackbone(nn.Module):
    """Hourglass backbone"""
    def __init__(self, name: str,
                 train_backbone: bool,
                 return_interm_layers: bool):
        super().__init__()
        PRETRAINPATH = "/home/ubuntu/tx/projects/neurvps-master/log/scannet/paper-result.pth.tar"
        self.num_channels = [64]
        self.hg = hg(planes=self.num_channels[0], depth=4, num_stacks=1, num_blocks=1)
        # param_dict = torch.load(PRETRAINPATH, map_location='cpu')['model_state_dict']
        # new_param_dict = OrderedDict()
        # for k, v in param_dict.items():
        #     if "module.backbone" in k:
        #         new_param_dict[k[16:]] = v
        # self.hg.load_state_dict(new_param_dict, strict=True)

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


def build_backbone(args):
    train_backbone = args.lr_backbone > 0
    return_interm_layers = args.return_interm_layers
    if "resnet" in args.backbone.lower():
        backbone = Backbone(args.backbone, train_backbone, return_interm_layers)
    elif "hourglass" in args.backbone.lower():
        backbone = HourglassBackbone(args.backbone, train_backbone, return_interm_layers)
    else:
        raise NotImplementedError
    return backbone
