import math
import megengine as mge
import megengine.functional as F
import numpy as np
from megengine import Tensor
import pdb

def restore_bbox(rois, deltas, unnormalize=True, config = None):

    assert deltas.ndim == 3
    if unnormalize:
        std_opr = mge.tensor(config.bbox_normalize_stds.reshape(1, 1, -1))
        mean_opr = mge.tensor(config.bbox_normalize_means.reshape(1, 1, -1))
        deltas = deltas * std_opr
        deltas = deltas + mean_opr

    # n = deltas.shape[1]
    n, c = deltas.shape[0], deltas.shape[1]
    all_rois = F.broadcast_to(F.expand_dims(rois, 1), (n, c, rois.shape[1])).reshape(-1, rois.shape[1])
    deltas = deltas.reshape(-1, deltas.shape[2])
    pred_bbox = bbox_transform_inv_opr(all_rois, deltas)
    pred_bbox = pred_bbox.reshape(-1, c, pred_bbox.shape[1])
    return pred_bbox

def filter_boxes_opr(boxes, min_size):

    """Remove all boxes with any side smaller than min_size."""
    wh = boxes[:, 2:4] - boxes[:, 0:2] + 1
    keep_mask = F.prod(wh >= min_size, axis = 1).astype(np.float32)
    keep_mask = keep_mask + F.equal(keep_mask.sum(), 0).astype(np.float32)
    return keep

def clip_boxes_opr(boxes, im_info):
    """ Clip the boxes into the image region."""
    w = im_info[1] - 1
    h = im_info[0] - 1
    boxes[:, 0::4] = boxes[:, 0::4].clamp(min=0, max=w)
    boxes[:, 1::4] = boxes[:, 1::4].clamp(min=0, max=h)
    boxes[:, 2::4] = boxes[:, 2::4].clamp(min=0, max=w)
    boxes[:, 3::4] = boxes[:, 3::4].clamp(min=0, max=h)
    return boxes

def bbox_transform_inv_opr(bbox, deltas):
    max_delta = math.log(1000.0 / 16)
    """ Transforms the learned deltas to the final bbox coordinates, the axis is 1"""
    bbox_width = bbox[:, 2] - bbox[:, 0] + 1
    bbox_height = bbox[:, 3] - bbox[:, 1] + 1
    bbox_ctr_x = bbox[:, 0] + 0.5 * bbox_width
    bbox_ctr_y = bbox[:, 1] + 0.5 * bbox_height
    pred_ctr_x = bbox_ctr_x + deltas[:, 0] * bbox_width
    pred_ctr_y = bbox_ctr_y + deltas[:, 1] * bbox_height

    dw = deltas[:, 2]
    dh = deltas[:, 3]
    dw = 
F.minimum(dw, max_delta)