import torch
import torch.nn as nn
import copy 
from .point_generator import MlvlPointGenerator
from .utils import NMS, bbox_decode


class YOLOXPostProcess(nn.Module):
    def __init__(self, classes=[1], num_classes=1, strides=[8, 16, 32], nms_params=None, cal_max_score=True):
        super(YOLOXPostProcess, self).__init__()
        self.classes = classes
        self.num_classes = num_classes

        self.prior_generator = MlvlPointGenerator(strides=strides, offset=0)

        self.pre_nms = nms_params.pop('pre_nms', 1000)
        self.nms = NMS(**nms_params)
        self.cal_max_score = cal_max_score

    def forward(self, inputs):
        # preds
        obj_pred_list, cls_pred_list, box_pred_list = inputs
            
        batch_size = cls_pred_list[0].size(0)
        assert batch_size == 1

        # priors
        featmap_sizes = [feat.shape[2:] for feat in cls_pred_list]     
 
        priors = self.prior_generator.get_priors(featmap_sizes,
                                       dtype=cls_pred_list[0].dtype,
                                       device=cls_pred_list[0].device,
                                       with_stride=True)
        flatten_priors = torch.cat(priors)

        cls_pred_list = [cls_pred.permute(0, 2, 3, 1).reshape((batch_size, -1, self.num_classes))
                         for cls_pred in cls_pred_list]
        box_pred_list = [box_pred.permute(0, 2, 3, 1).reshape((batch_size, -1, 4))
                         for box_pred in box_pred_list]
        obj_pred_list = [box_pred.permute(0, 2, 3, 1).reshape((batch_size, -1, 1))
                         for box_pred in obj_pred_list]

        cls_pred = torch.cat(cls_pred_list, dim=1).sigmoid()  # (N, HWA, C)
        box_pred = torch.cat(box_pred_list, dim=1)  # (N, HWA, 4)
        obj_pred = torch.cat(obj_pred_list, dim=1).sigmoid()


        decode_box = bbox_decode(flatten_priors, box_pred)

        # post process
        pred_boxes_list = []
        pred_classes_list = []
        pred_scores_list = []
        for i in range(batch_size):
            pred_boxes, pred_classes, pred_scores = self._process_single_image(cls_pred[i], obj_pred[i], decode_box[i])
            pred_boxes_list.append(pred_boxes.detach().cpu())
            pred_classes_list.append(pred_classes.detach().cpu())
            pred_scores_list.append(pred_scores.detach().cpu())
        return pred_boxes_list[0], pred_classes_list[0], pred_scores_list[0]

    def _process_single_image(self, cls_preds, obj_preds, box_preds):
        real_classes = cls_preds[0].new_tensor(self.classes, dtype=torch.long)
        pred_boxes = []
        pred_scores = []
        pred_classes = []
        for cls_pred, obj_pred, box_pred in zip([cls_preds], [obj_preds], [box_preds]):
            score = cls_pred * obj_pred
            if self.cal_max_score:
                pred_score, pred_class = score.max(1)
                pred_class = real_classes[pred_class]
                if self.pre_nms > 0:
                    num_topk = min(pred_score.shape[0], self.pre_nms)
                else:
                    num_topk = pred_score.shape[0]
                pred_score, topk_ind = pred_score.sort(descending=True)
                pred_score = pred_score[:num_topk]
                topk_ind = topk_ind[:num_topk]
                pred_class = pred_class[topk_ind]
                box_pred = box_pred[topk_ind]
            else:
                pred_score = score.reshape(-1)
                if self.pre_nms > 0:
                    num_topk = min(pred_score.shape[0], self.pre_nms)
                else:
                    num_topk = pred_score.shape[0]
                pred_score, topk_ind = pred_score.sort(descending=True)
                pred_score = pred_score[:num_topk]
                topk_ind = topk_ind[:num_topk]
                pred_class = topk_ind % self.num_classes
                topk_ind = topk_ind // self.num_classes
                pred_class = real_classes[pred_class]
                box_pred = box_pred[topk_ind]
            pred_boxes.append(box_pred)
            pred_scores.append(pred_score)
            pred_classes.append(pred_class)
        pred_boxes = torch.cat(pred_boxes, dim=0)
        pred_scores = torch.cat(pred_scores, dim=0)
        pred_classes = torch.cat(pred_classes, dim=0)
        if self.nms is not None:
            pred_dets = torch.cat([pred_boxes, pred_scores.view(-1, 1), pred_classes.view(-1, 1)], dim=1)
            pred_dets = self.nms(pred_dets)
            pred_boxes = pred_dets[:, :-2]
            pred_scores = pred_dets[:, -2]
            pred_classes = pred_dets[:, -1].to(torch.long)
        return pred_boxes, pred_classes, pred_scores
