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)