#
# Modified by Peize Sun
# Contact: sunpeize@foxmail.com
#
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
OneNet Transformer class.

Copy-paste from torch.nn.Transformer with modifications:
    * positional encodings are passed in MHattention
    * extra LN at the end of encoder is removed
    * decoder returns a stack of activations from all decoding layers
"""
import copy
import math
from typing import Optional, List

import torch
from torch import nn, Tensor
import torch.nn.functional as F

from detectron2.modeling.poolers import ROIPooler, cat
from detectron2.structures import Boxes
from .deconv import CenternetDeconv

class Head(nn.Module):

    def __init__(self, cfg, backbone_shape=[2048, 1024, 512, 256]):
        super().__init__()
        
        # Build heads.
        num_classes = cfg.MODEL.OneNet.NUM_CLASSES
        d_model = cfg.MODEL.OneNet.DECONV_CHANNEL[-1]
        activation = cfg.MODEL.OneNet.ACTIVATION

        self.deconv = CenternetDeconv(cfg, backbone_shape)
        
        self.num_classes = num_classes
        self.d_model = d_model
        self.num_classes = num_classes
        self.activation = _get_activation_fn(activation)

        self.feat1 = nn.Conv2d(self.d_model, self.d_model, kernel_size=3, stride=1, padding=1)
        self.cls_score = nn.Conv2d(d_model, num_classes, kernel_size=3, stride=1, padding=1)
        self.ltrb_pred = nn.Conv2d(d_model, 4, kernel_size=3, stride=1, padding=1)        
        
        # Init parameters.
        prior_prob = cfg.MODEL.OneNet.PRIOR_PROB
        self.bias_value = -math.log((1 - prior_prob) / prior_prob)
        self._reset_parameters()

    def _reset_parameters(self):
        # init all parameters.
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

            # initialize the bias for focal loss.
            if p.shape[-1] == self.num_classes:
                nn.init.constant_(p, self.bias_value)
    
    def forward(self, features_list):
        
        features = self.deconv(features_list)
        locations = self.locations(features)[None]       

        feat = self.activation(self.feat1(features))
    
        class_logits = self.cls_score(feat)
        pred_ltrb = F.relu(self.ltrb_pred(feat))
        pred_bboxes = self.apply_ltrb(locations, pred_ltrb)

        return class_logits, pred_bboxes
    
    def apply_ltrb(self, locations, pred_ltrb): 
        """
        :param locations:  (1, 2, H, W)
        :param pred_ltrb:  (N, 4, H, W) 
        """

        pred_boxes = torch.zeros_like(pred_ltrb)
        pred_boxes[:,0,:,:] = locations[:,0,:,:] - pred_ltrb[:,0,:,:]  # x1
        pred_boxes[:,1,:,:] = locations[:,1,:,:] - pred_ltrb[:,1,:,:]  # y1
        pred_boxes[:,2,:,:] = locations[:,0,:,:] + pred_ltrb[:,2,:,:]  # x2
        pred_boxes[:,3,:,:] = locations[:,1,:,:] + pred_ltrb[:,3,:,:]  # y2

        return pred_boxes    
    
    @torch.no_grad()
    def locations(self, features, stride=4):
        """
        Arguments:
            features:  (N, C, H, W)
        Return:
            locations:  (2, H, W)
        """

        h, w = features.size()[-2:]
        device = features.device
        
        shifts_x = torch.arange(
            0, w * stride, step=stride,
            dtype=torch.float32, device=device
        )
        shifts_y = torch.arange(
            0, h * stride, step=stride,
            dtype=torch.float32, device=device
        )
        shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
        shift_x = shift_x.reshape(-1)
        shift_y = shift_y.reshape(-1)
        locations = torch.stack((shift_x, shift_y), dim=1) + stride // 2            
        
        locations = locations.reshape(h, w, 2).permute(2, 0, 1)
        
        return locations


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")

[((89, 5, 89, 20), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((43, 21, 43, 94), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', False, 'from torch import nn, Tensor\n'), ((44, 25, 44, 92), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', False, 'from torch import nn, Tensor\n'), ((45, 25, 45, 82), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', False, 'from torch import nn, Tensor\n'), ((81, 21, 81, 48), 'torch.zeros_like', 'torch.zeros_like', ({(81, 38, 81, 47): 'pred_ltrb'}, {}), '(pred_ltrb)', False, 'import torch\n'), ((101, 19, 104, 9), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((105, 19, 108, 9), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((109, 27, 109, 61), 'torch.meshgrid', 'torch.meshgrid', ({(109, 42, 109, 50): 'shifts_y', (109, 52, 109, 60): 'shifts_x'}, {}), '(shifts_y, shifts_x)', False, 'import torch\n'), ((49, 27, 49, 66), 'math.log', 'math.log', ({(49, 36, 49, 65): '((1 - prior_prob) / prior_prob)'}, {}), '((1 - prior_prob) / prior_prob)', False, 'import math\n'), ((112, 20, 112, 58), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((56, 16, 56, 42), 'torch.nn.init.xavier_uniform_', 'nn.init.xavier_uniform_', ({(56, 40, 56, 41): 'p'}, {}), '(p)', False, 'from torch import nn, Tensor\n'), ((60, 16, 60, 53), 'torch.nn.init.constant_', 'nn.init.constant_', ({(60, 34, 60, 35): 'p', (60, 37, 60, 52): 'self.bias_value'}, {}), '(p, self.bias_value)', False, 'from torch import nn, Tensor\n')]