import torch
import torch.nn as nn
from torchvision.ops.boxes import batched_nms

from util import box_ops


class CondNMSPostProcess(nn.Module):
    def __init__(self, num_queries):
        super(CondNMSPostProcess, self).__init__()
        self.num_queries = num_queries

    @torch.no_grad()
    def forward(self, outputs, target_sizes, pred_names, mask_infos):
        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']
        bs = len(out_logits)
        results = []

        for b in range(bs):
            b_scores, b_boxes, b_names = [], [], []
            b_start_id, b_end_id = [], []
            name = []
            for name_i in pred_names[b]:
                name.append([name_i] * self.num_queries)
            start_id, end_id = [], []
            for (start, end) in mask_infos[b].keys():
                start_id.append([start] * self.num_queries)
                end_id.append([end] * self.num_queries)
            prob = out_logits[b][0][:, -1:].sigmoid()
            if len(prob) == 0:
                continue
            boxes = box_ops.box_cxcywh_to_xyxy(out_bbox[b][0])
            img_h, img_w = target_sizes[b]
            scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=0)
            boxes = boxes * scale_fct[None, :]
            num_patch = len(prob) // self.num_queries
            prob = prob.view(num_patch, self.num_queries, -1)
            boxes = boxes.view(num_patch, self.num_queries, -1)
            for t in range(num_patch):
                ind = prob[t].squeeze(1).topk(100).indices
                prob_prenms = prob[t][ind]
                box_prenms = boxes[t][ind]
                lbl_prenms = torch.zeros_like(prob_prenms)
                nms_ind = batched_nms(box_prenms, prob_prenms[:, 0], lbl_prenms[:, 0], 0.7)[:20]
                b_scores.append(prob_prenms[nms_ind])
                b_boxes.append(box_prenms[nms_ind])

                b_names += [name[t][int(i)] for i in nms_ind]
                b_start_id += [start_id[t][int(i)] for i in nms_ind]
                b_end_id += [end_id[t][int(i)] for i in nms_ind]
            b_scores = torch.cat(b_scores).cpu().squeeze(1)
            b_boxes = torch.cat(b_boxes).cpu()
            out = {'scores': b_scores, 'boxes': b_boxes, 'names': b_names,
                   'start_id': b_start_id, 'end_id': b_end_id}
            results.append(out)
        return results