import numpy as np
import megengine as mge
import megengine.functional as F
import megengine.module as M
from config import config
from backbone.resnet50 import ResNet50
from module.rpn import RPN
from layers.roi_pool import roi_pool
from det_opr.bbox_opr import bbox_transform_inv_opr, restore_bbox
from det_opr.fpn_roi_target import fpn_roi_target
from det_opr.loss_opr import softmax_loss_opr, smooth_l1_loss_rcnn_opr
from det_opr.utils import get_padded_tensor
import pdb
class Network(M.Module):
    def __init__(self):
        super().__init__()
        # ----------------------- build the backbone ------------------------ #
        self.resnet50 = ResNet50()
        # ------------ freeze the weights of resnet stage1 and stage 2 ------ #
        if config.backbone_freeze_at >= 1:
            for p in self.resnet50.conv1.parameters():
                # p.requires_grad = False
                p = p.detach()
        if config.backbone_freeze_at >= 2:
            for p in self.resnet50.layer1.parameters():
                # p.requires_grad = False
                p = p.detach()
        # -------------------------- build the FPN -------------------------- #
        self.backbone = FPN(self.resnet50)
        # -------------------------- build the RPN -------------------------- #
        self.RPN = RPN(config.rpn_channel)
        # ----------------------- build the RCNN head ----------------------- #
        self.RCNN = RCNN()
        # -------------------------- input Tensor --------------------------- #
        self.inputs = {
            "image": mge.tensor(
                np.random.random([2, 3, 224, 224]).astype(np.float32), dtype="float32",
            ),
            "im_info": mge.tensor(
                np.random.random([2, 5]).astype(np.float32), dtype="float32",
            ),
            "gt_boxes": mge.tensor(
                np.random.random([2, 100, 5]).astype(np.float32), dtype="float32",
            ),
        }
    
    def pre_process(self, images):

        mean = config.image_mean.reshape(1, 3, 1, 1).astype(np.float32)
        std = config.image_std.reshape(1, 3, 1, 1).astype(np.float32)
        mean = mge.tensor(mean).to(images.device)
        std = mge.tensor(std).to(images.device)
        normed_images = (images - mean) / std
        normed_images = get_padded_tensor(normed_images, 64)
        return normed_images

    def forward(self, inputs):

        images = inputs['image']
        im_info = inputs['im_info']
        gt_boxes = inputs['gt_boxes']
        #del images
        # process the images
        normed_images = self.pre_process(images)
        if self.training:
            return self._forward_train(normed_images, im_info, gt_boxes)
        else:
            return self._forward_test(normed_images, im_info)

    def _forward_train(self, image, im_info, gt_boxes):

        loss_dict = {}
        # stride: 64,32,16,8,4, p6->p2
        fpn_fms = self.backbone(image)
        rpn_rois, loss_dict_rpn = \
            self.RPN(fpn_fms, im_info, gt_boxes)
        rcnn_rois, rcnn_labels, rcnn_bbox_targets = fpn_roi_target(
            rpn_rois, im_info, gt_boxes, top_k=2)
        loss_dict_rcnn = self.RCNN(
                fpn_fms, rcnn_rois, rcnn_labels, rcnn_bbox_targets)
        loss_dict.update(loss_dict_rpn)
        loss_dict.update(loss_dict_rcnn)
        return loss_dict

    def _forward_test(self, image, im_info):

        fpn_fms = self.backbone(image)
        
        rpn_rois = self.RPN(fpn_fms, im_info)
        
        pred_bbox = self.RCNN(fpn_fms, rpn_rois)
        return pred_bbox

class RCNN(M.Module):
    def __init__(self):
        super().__init__()
        # roi head
        self.refinement = True
        self.fc1 = M.Linear(256*7*7, 1024)
        self.fc2 = M.Linear(1024, 1024)
        self.fc3 = M.Linear(1054, 1024) if self.refinement else None

        self.relu = M.ReLU()

        self.n = config.num_classes
        self.a = M.Linear(1024, 5 * self.n)
        self.b = M.Linear(1024, 5 * self.n)

        self.q = M.Linear(1024, 5 * self.n) if self.refinement else None
        self.r = M.Linear(1024, 5 * self.n) if self.refinement else None
        self._init_weights()

    def _init_weights(self,):
        
        for l in [self.fc1, self.fc2, self.a, self.b]:
            M.init.normal_(l.weight, std=0.01)
            
M.init.fill_(l.bias, 0)