# -*- coding: utf-8 -*-
import megengine as mge
import megengine.random as rand
import megengine.functional as F
import numpy as np
from config import config
from det_opr.bbox_opr import box_overlap_opr, bbox_transform_opr, box_overlap_ignore_opr
import pdb

def fpn_roi_target(rpn_rois, im_info, gt_boxes, fg_threshold = config.fg_threshold, top_k=1):

    return_rois, return_labels = [], []
    return_bbox_targets = []
    # get per image proposals and gt_boxes
    batch_per_gpu = im_info.shape[0]
    sampling = True 
    # is_sample = True if top_k < 2 else False
    for bid in range(batch_per_gpu):

        gt_boxes_perimg = gt_boxes[bid, :im_info[bid, 5].astype(np.int32), :]
        dummy_gt = F.ones([1, gt_boxes_perimg.shape[1]])

        batch_inds = F.ones((gt_boxes_perimg.shape[0], 1)) * bid
        #if config.proposal_append_gt:
        gt_rois = F.concat([batch_inds, gt_boxes_perimg[:, :4]], axis=1)
        batch_rois_mask = F.equal(rpn_rois[:, 0], bid) > 0
        _, batch_rois_index = F.cond_take(batch_rois_mask, batch_rois_mask)
        
        # batch_roi_mask = rpn_rois[:, 0] == bid
        # batch_roi_inds = mask_to_inds(batch_roi_mask)
        all_rois= F.concat([rpn_rois[batch_rois_index], gt_rois], axis=0) if sampling \
            else rpn_rois[batch_rois_index]
        # all_rois = F.concat([rpn_rois.ai[batch_roi_inds], gt_rois], axis=0)

        gt_boxes_perimg = F.concat([gt_boxes_perimg, dummy_gt],axis=0)
        overlaps_normal, overlaps_ignore = box_overlap_ignore_opr(
                all_rois[:, 1:5], gt_boxes_perimg)

        # overlaps_normal, overlaps_normal_indices = F.argsort(overlaps_normal, descending=True)
        # overlaps_ignore, overlaps_ignore_indices = F.argsort(overlaps_ignore, descending=True)
        overlaps_normal_indices = F.argsort(overlaps_normal, descending=True)
        overlaps_normal = F.gather(overlaps_normal, 1, overlaps_normal_indices)
        # overlaps_normal = F.nn.indexing_one_hot(overlaps_normal, overlaps_normal_indices, 1)
        overlaps_ignore_indices = F.argsort(overlaps_ignore, descending = True)
        overlaps_ignore = F.gather(overlaps_ignore, 1, overlaps_ignore_indices)
        # overlaps_ignore = F.nn.indexing_one_hot(overlaps_ignore, overlaps_ignore_indices, 1)


        # gt max and indices, ignore max and indices
        max_overlaps_normal = overlaps_normal[:, :top_k].flatten()
        gt_assignment_normal = overlaps_normal_indices[:, :top_k].flatten()
        max_overlaps_ignore = overlaps_ignore[:, :top_k].flatten()
        gt_assignment_ignore = overlaps_ignore_indices[:, :top_k].flatten()
        # cons masks
        
        ignore_assign_mask = (max_overlaps_normal < fg_threshold).astype(np.float32) * (
                max_overlaps_ignore > max_overlaps_normal).astype(np.float32)
        max_overlaps = max_overlaps_normal * (1 - ignore_assign_mask).astype(np.float32) + \
                max_overlaps_ignore * ignore_assign_mask
        

        gt_assignment = gt_assignment_normal * (1- ignore_assign_mask) + \
                gt_assignment_ignore * ignore_assign_mask
        
        gt_assignment = gt_assignment.astype(np.int32)

        labels = gt_boxes_perimg[gt_assignment, 4]
        fg_mask = (max_overlaps >= fg_threshold).astype(np.float32) * (1 - F.equal(labels, config.ignore_label))
        bg_mask = (max_overlaps < config.bg_threshold_high).astype(np.float32) * (
                max_overlaps >= config.bg_threshold_low).astype(np.float32)

        fg_mask = fg_mask.reshape(-1, top_k)
        bg_mask = bg_mask.reshape(-1, top_k)
        pos_max = config.num_rois * config.fg_ratio
        fg_inds_mask = _bernoulli_sample_masks(fg_mask[:, 0], pos_max, 1) if sampling else 
F.equal(fg_mask[:, 0], 0)