import os, sys
import numpy as np
from config import config
from det_opr.bbox_opr import box_overlap_opr, bbox_transform_opr
import megengine as mge
from megengine import functional as F
import pdb
def _compute_center(boxes):

    ptrx = 0.5 * (boxes[:, 0] + boxes[:, 2])
    ptry = 0.5 * (boxes[:, 1] + boxes[:, 3])
    centre = F.stack([ptrx, ptry], axis=1)
    return centre

def _compute_pos_area(gtboxes, ratio = 0.3):

    H, W = gtboxes[:, 3] - gtboxes[:, 1], gtboxes[:, 2] - gtboxes[:, 0]
    centres = _compute_center(gtboxes)
    l = centres[:, 0] - ratio * W
    r = centres[:, 0] + ratio * W
    t = centres[:, 1] - ratio * H
    b = centres[:, 1] + ratio * H
    boundary = F.stack([l, t, r, b], axis = 1)
    return boundary

def _anchor_double_target(gt_boxes, im_info, all_anchors):

    gt_boxes, im_info = gt_boxes.detach(), im_info.detach()
    all_anchors = all_anchors.detach()
    
    gt_boxes = gt_boxes[:im_info[5].astype(np.int32), :]
    dummy = -F.ones([1, gt_boxes.shape[1]]).to(gt_boxes.device)
    gt_boxes = F.concat([gt_boxes, dummy], axis=0)
    valid_mask = 1 - (gt_boxes[:, 4] < 0).astype(np.float32)

    anchor_centers = _compute_center(all_anchors)
    gtboxes_centers = _compute_center(gt_boxes)
    # gtboxes_centers = gtboxes_centers * valid_mask.unsqueeze(1)
    gtboxes_centers = gtboxes_centers * F.expand_dims(valid_mask, axis=1)

    N, K = all_anchors.shape[0], gt_boxes.shape[0]
    an_centers = F.expand_dims(anchor_centers, axis=1)
    gt_centers = F.expand_dims(gtboxes_centers, axis=0)
    # an_centers = anchor_centers.unsqueeze(1).repeat(1, K, 1)
    # gt_centers = gtboxes_centers.unsqueeze(0).repeat(N, 1, 1)

    distance = F.abs(an_centers - gt_centers)
    distance = F.sqrt(F.pow(distance, 2).sum(axis=2))
    
    start = 0
    end = 5
    overlaps = box_overlap_opr(all_anchors[:, :4], gt_boxes[:, :4])
    overlaps *= F.expand_dims(valid_mask, axis=0)
    default_num = 16

    ious_list = []

    for l in range(start, end):

        _, index = F.cond_take(all_anchors[:, 4] == l, all_anchors[:, 4])

        level_dist = distance[index, :].transpose(1, 0)
        ious = overlaps[index, :].transpose(1, 0)
        sorted_index = F.argsort(level_dist, descending=False)
        n = min(sorted_index.shape[1], default_num)
        ious = 
F.gather(ious, 1, sorted_index[:, :n])