import numpy as np
import megengine as mge
import megengine.functional as F
import megengine.module as M
import math
from config import config
from backbone.resnet50 import ResNet50
from module.generate_anchors import generate_anchors
from det_opr.bbox_opr import bbox_transform_inv_opr, box_overlap_opr
from det_opr.utils import get_padded_tensor
from rpn_anchor_target_opr import rpn_anchor_target_opr
from det_opr.loss_opr import sigmoid_cross_entropy_retina, smooth_l1_loss_retina, iou_l1_loss
import pdb

class RetinaNetAnchorV2(M.Module):
    def __init__(self):
        super().__init__()
    
    def generate_anchors_opr(self, fm_3x3, fm_stride,
        anchor_scales=(8, 16, 32, 64, 128), 
        anchor_ratios=(1, 2, 3), base_size = 4):

        np_anchors = generate_anchors(
            base_size=base_size,
            ratios=np.array(anchor_ratios),
            scales=np.array(anchor_scales))
        device = fm_3x3.device
        anchors = mge.tensor(np_anchors).to(device)
        height, width = fm_3x3.shape[2], fm_3x3.shape[3]
        shift_x = F.linspace(0, width-1, width).to(device) * fm_stride
        shift_y = F.linspace(0, height -1, height).to(device) * fm_stride

        broad_shift_x = F.broadcast_to(shift_x.reshape(1, -1), (height, width)).flatten()
        broad_shift_y = F.broadcast_to(shift_y.reshape(-1, 1), (height, width)).flatten()
        shifts = F.stack([broad_shift_x, broad_shift_y, broad_shift_x, broad_shift_y], axis=1)

        c = anchors.shape[1]
        all_anchors = F.expand_dims(anchors, axis=0) + F.expand_dims(shifts, axis=1)
        all_anchors = all_anchors.reshape(-1, c).detach()
        return all_anchors

    def forward(self, fpn_fms):

        all_anchors_list = []
        fm_stride = [8, 16, 32, 64, 128]
        fm_stride.reverse()

        for i, fm_3x3 in enumerate(fpn_fms):
            
            anchor_scales = np.array(config.anchor_base_scale) * fm_stride[i]
            all_anchors = self.generate_anchors_opr(fm_3x3, fm_stride[i], anchor_scales,
                config.anchor_aspect_ratios, base_size = 4)
            all_anchors_list.append(all_anchors)
        return all_anchors_list

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)
        self.head = RetinaNetHead()
        # -------------------------- buid the anchor generator -------------- #
        self.anchor_generator = RetinaNetAnchorV2()
        
        # -------------------------- buid the criteria ---------------------- #
        self.criteria = RetinaNetCriteriaV2()
        # -------------------------- input Tensor --------------------------- #
        self.inputs = {
            "image": mge.tensor(
                np.random.random([2, 3, 756, 1400]).astype(np.float32), dtype="float32",
            ),
            "im_info": mge.tensor(
                np.random.random([2, 6]).astype(np.float32), dtype="float32",
            ),
            "gt_boxes": mge.tensor(
                np.random.random([2, 500, 5]).astype(np.float32), dtype="float32",
            ),
        }
    
    def pre_process(self, images):

        mean = config.image_mean.reshape(1, -1, 1, 1).astype(np.float32)
        std = config.image_std.reshape(1, -1, 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):

        im_info = inputs['im_info']
        # process the images
        normed_images = self.pre_process(inputs['image'])

        if self.training:
            gt_boxes = inputs['gt_boxes']
            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: 128,64,32,16,8, p6->p2
        fpn_fms = self.backbone(image)
        pred_cls_list, rpn_num_prob_list, pred_reg_list, rpn_iou_list = self.head(fpn_fms)

        anchors_list = self.anchor_generator(fpn_fms)

        loss_dict = self.criteria(
                pred_cls_list, rpn_num_prob_list, pred_reg_list, anchors_list, 
                rpn_iou_list, gt_boxes, im_info)

        return loss_dict

    def _forward_test(self, image, im_info):

        fpn_fms = self.backbone(image)

        pred_cls_list, rpn_num_prob_list, pred_reg_list, rpn_iou_list = self.head(fpn_fms)

        anchors_list = self.anchor_generator(fpn_fms)

        pred_boxes = self._recover_dtboxes(anchors_list, pred_cls_list,
            pred_reg_list, rpn_iou_list)

        return pred_boxes

    def _recover_dtboxes(self, anchors_list, rpn_cls_list, rpn_bbox_list, rpn_iou_list):

        assert rpn_cls_list[0].shape[0] == 1
        all_anchors = 
F.concat(anchors_list, axis = 