"""
Grid-feat-vqa
"""
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.utils.logger import setup_logger
from detectron2.modeling import build_model
from detectron2.layers import FrozenBatchNorm2d
import torch
from torch import nn
from src.modeling.grid_feats import add_attribute_config
import os
import horovod.torch as hvd


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1,
            init_identity=False):
    """3x3 convolution with padding"""
    conv = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False,
                     dilation=dilation)
    if init_identity:
        """ init as identity kernel, not working yet
        0 0 0
        0 1 0
        0 0 0
        """
        identity_weight = conv.weight.new_zeros(3, 3)
        identity_weight[1, 1] = 1. / in_planes
        identity_weight = identity_weight.view(
            1, 1, 3, 3).expand(conv.weight.size())
        with torch.no_grad():
            conv.weight = nn.Parameter(identity_weight)
    return conv


class GridFeatBackbone(nn.Module):
    def __init__(self, detectron2_model_cfg, config,
                 input_format="BGR"):
        super(GridFeatBackbone, self).__init__()
        self.detectron2_cfg = self.__setup__(detectron2_model_cfg)
        self.feature = build_model(self.detectron2_cfg)
        self.grid_encoder = nn.Sequential(
            conv3x3(config.backbone_channel_in_size,
                    config.hidden_size),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
        )
        self.input_format = input_format
        assert input_format == "BGR", "detectron 2 image input format should be BGR"
        
        self.config = config

    def __setup__(self, config_file):
        """
        Create configs and perform basic setups.
        """
        rank = hvd.rank()
        detectron2_cfg = get_cfg()
        add_attribute_config(detectron2_cfg)
        detectron2_cfg.merge_from_file(config_file)
        # force the final residual block to have dilations 1
        detectron2_cfg.MODEL.RESNETS.RES5_DILATION = 1
        # FIXME: hardcode device to each rank
        detectron2_cfg.MODEL.DEVICE = f'cpu'
        detectron2_cfg.freeze()

        setup_logger(None, distributed_rank=rank, name="fvcore")
        logger = setup_logger(None, distributed_rank=rank)
        return detectron2_cfg

    def load_state_dict(self, model_path):
        if not os.path.exists(model_path):
            print(f"{model_path} does not exist, loading ckpt from detectron2")
            DetectionCheckpointer(
                self.feature).resume_or_load(
                    self.detectron2_cfg.MODEL.WEIGHTS, resume=True)
        else:
            DetectionCheckpointer(self.feature).resume_or_load(
                    model_path, resume=True)

    @property
    def config_file(self):
        return self.detectron2_cfg.dump()

    def train(self, mode=True):
        super(GridFeatBackbone, self).train(mode)

    def forward(self, x):
        bsz, n_frms, c, h, w = x.shape
        x = x.view(bsz*n_frms, c, h, w)
        if self.input_format == "BGR":
            # RGB->BGR, images are read in as RGB by default
            x = x[:, [2, 1, 0], :, :]
        res5_features = self.feature.backbone(x)
        grid_feat_outputs = self.feature.roi_heads.get_conv5_features(
            res5_features)

        grid = self.grid_encoder(grid_feat_outputs)  # (B * n_frm, C, H, W)
        new_c, new_h, new_w = grid.shape[-3:]
        # if n_frms != 0:
        grid = grid.view(bsz, n_frms, new_c, new_h, new_w)  # (B, n_frm, C, H, W)

        grid = grid.permute(0, 1, 3, 4, 2)  # (B, n_frm=3, H, W, C)
        return grid
