import torch
import torch.nn as nn
import numpy as np
from torch.nn import functional as F

def get_width(boxes, offset=0):
    return boxes[:, 2] - boxes[:, 0] + offset


def get_height(boxes, offset=0):
    return boxes[:, 3] - boxes[:, 1] + offset


def get_size(boxes, offset=0):
    return get_width(boxes, offset) * get_height(boxes, offset)


def overlaps(boxes_a, boxes_b, mode='iou', is_aligned=False, offset=0, eps=1e-6):
    # boxes_a: (n, 4)
    # boxes_b: (m, 4)
    # return: (n, m)
    if is_aligned:
        assert boxes_a.shape[0] == boxes_b.shape[0]
    if boxes_a.shape[0] * boxes_b.shape[0] == 0:
        if is_aligned:
            if isinstance(boxes_a, torch.Tensor):
                return boxes_a.new_zeros((boxes_a.shape[0], ))
            else:
                return np.zeros((boxes_a.shape[0], ))
        else:
            if isinstance(boxes_a, torch.Tensor):
                return boxes_a.new_zeros((boxes_a.shape[0], boxes_b.shape[0]))
            else:
                return np.zeros((boxes_a.shape[0], boxes_b.shape[0]))
    if is_aligned:
        if isinstance(boxes_a, torch.Tensor):
            lt = torch.max(boxes_a[:, :2], boxes_b[:, :2])  # (n, 2)
            rb = torch.min(boxes_a[:, 2:4], boxes_b[:, 2:4])  # (n, 2)
            wh = (rb - lt + offset).clamp(min=0)  # (n, 2)
        else:
            lt = np.maximum(boxes_a[:, :2], boxes_b[:, :2])  # (n, 2)
            rb = np.minimum(boxes_a[:, 2:4], boxes_b[:, 2:4])  # (n, 2)
            wh = np.maximum(rb - lt + offset, 0)  # (n, 2)
        overlap = wh[:, 0] * wh[:, 1]  # (n,)
        if mode is None:
            return overlap
        area_a = get_size(boxes_a, offset=offset)  # (n,)
        area_b = get_size(boxes_b, offset=offset)  # (n,)
    # intersection over union
        union = area_a + area_b - overlap  # (n,)
    else:
        if isinstance(boxes_a, torch.Tensor):
            lt = torch.max(boxes_a[:, None, :2], boxes_b[None, :, :2])  # (n, m, 2)
            rb = torch.min(boxes_a[:, None, 2:4], boxes_b[None, :, 2:4])  # (n, m, 2)
            wh = (rb - lt + offset).clamp(min=0)  # (n, m, 2)
        else:
            lt = np.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])  # (n, m, 2)
            rb = np.minimum(boxes_a[:, None, 2:4], boxes_b[None, :, 2:4])  # (n, m, 2)
            wh = np.maximum(rb - lt + offset, 0)  # (n, m, 2)
        overlap = wh[:, :, 0] * wh[:, :, 1]  # (n, m)
        if mode is None:
            return overlap
        area_a = get_size(boxes_a, offset=offset)  # (n, )
        area_b = get_size(boxes_b, offset=offset)  # (m, )
    # intersection over union
        union = area_a[:, None] + area_b[None, :] - overlap  # (n, m)

    if isinstance(boxes_a, torch.Tensor):
        union = union.clamp(min=eps)
    else:
        union = np.maximum(union, eps)
    ios = overlap / union
    return ios

class Overlap2D(nn.Module):
    def __init__(self, mode='iou', is_aligned=False, offset=0, eps=1e-6):
        super(Overlap2D, self).__init__()
        self.mode = mode
        self.is_aligned = is_aligned
        self.offset = offset
        self.eps = eps

    def extra_repr(self):
        s = 'mode={}'.format(self.mode)
        s += ', is_aligned={}'.format(self.is_aligned)
        s += ', offset={}'.format(self.offset)
        s += ', eps={}'.format(self.eps)
        return s

    def forward(self, boxes_a, boxes_b, mode=None):
        mode = mode if mode is not None else self.mode
        return overlaps(boxes_a,
                        boxes_b,
                        mode=mode,
                        is_aligned=self.is_aligned,
                        offset=self.offset,
                        eps=self.eps)

class NMS(nn.Module):
    def __init__(self, nms_thresh, score_thresh=None, score_idx=None,
                 class_agnostic=False, pre_nms=-1, post_nms=-1,
                 overlap_params=None, method='nms', rescore_type='gaussian',
                 record_indices=False, record_suppressed_indices=False):
        super(NMS, self).__init__()
        assert method in ('nms', 'soft_nms')
        if overlap_params is None:
            overlap_params = {}
        if record_suppressed_indices:
            assert method == 'nms'
        self.nms_thresh = nms_thresh
        self.score_thresh = score_thresh
        self.score_idx = score_idx
        self.class_agnostic = class_agnostic
        self.pre_nms = pre_nms
        self.post_nms = post_nms
        self.overlap = Overlap2D(**overlap_params)
        self.method = method
        self.rescore_type = rescore_type
        self.record_indices = record_indices
        self.record_suppressed_indices = record_suppressed_indices

    def extra_repr(self):
        s = 'nms_thresh={}'.format(self.nms_thresh)
        if self.score_thresh is not None:
            s += ', score_thresh={}'.format(self.score_thresh)
        if self.score_idx is not None:
            s += ', score_idx={}'.format(self.score_idx)
        s += ', class_agnostic={}'.format(self.class_agnostic)
        if self.pre_nms > 0:
            s += ', pre_nms={}'.format(self.pre_nms)
        if self.post_nms > 0:
            s += ', post_nms={}'.format(self.post_nms)
        s += ', method={}'.format(self.method)
        if self.method == 'soft_nms':
            s += ', rescore_type={}'.format(self.rescore_type)
        if self.record_indices:
            s += ', record_indices={}'.format(self.record_indices)
        if self.record_suppressed_indices:
            s += ', record_suppressed_indices={}'.format(self.record_suppressed_indices)
        return s

    def forward(self, dets):
        # dets: (num_boxes, det_dim)
        det_dim = dets.size(1)
        post_nms = self.post_nms if self.post_nms > 0 else dets.size(0)
        score_idx = self.score_idx if self.score_idx is not None else det_dim - 2
        class_agnostic = self.class_agnostic or (det_dim - 1 == score_idx)
        # pad indices
        dets = F.pad(dets, (0, 1), mode='constant', value=0)
        dets[:, -1] = torch.arange(dets.size(0))
        # filter by score_thresh
        if self.score_thresh is not None:
            dets = dets[dets[:, score_idx] > self.score_thresh]
        if dets.numel() == 0:
            if not self.record_indices:
                dets = dets[:, :-1]
            if self.record_suppressed_indices:
                return dets, []
            else:
                return dets
        # sort by scores
        dets = dets[torch.argsort(dets[:, score_idx], descending=True)]
        if self.pre_nms > 0:
            dets = dets[:self.pre_nms]
        scores = dets[:, score_idx]
        order = torch.arange(dets.size(0))
        keep = []
        suppressed_indices = []
        while order.size(0) > 0 and len(keep) < post_nms:
            i = order[0]
            keep.append(i)
            iou = self.overlap(dets[i, None, :score_idx], dets[order[1:], :score_idx])[0]
            if not class_agnostic:
                iou = iou * (dets[i, score_idx+1] == dets[order[1:], score_idx+1]).float()
            if self.method == 'nms':
                if self.record_suppressed_indices:
                    inds = torch.nonzero(iou > self.nms_thresh, as_tuple=False).view(-1)
                    suppressed_indices.append(dets[order[inds + 1], -1])
                inds = torch.nonzero(iou <= self.nms_thresh, as_tuple=False).view(-1)
                order = order[inds + 1]
            elif self.method == 'soft_nms':
                # keep all boxes, but recalculate scores
                dets[i, score_idx] = scores[0]
                order = order[1:]
                scores = scores[1:]
                if self.rescore_type == 'linear':
                    inds = torch.nonzero(iou > self.nms_thresh, as_tuple=False).view(-1)
                    scores[inds] *= (1 - iou[inds])
                elif self.rescore_type == 'gaussian':
                    scores *= torch.exp(- iou ** 2 / self.nms_thresh)
                else:
                    assert False
                tmp = torch.argsort(scores, descending=True)
                order = order[tmp]
                scores = scores[tmp]
            else:
                assert False
        keep = dets.new_tensor(keep, dtype=torch.long)
        dets = dets[keep]
        if not self.record_indices:
            dets = dets[:, :-1]
        if self.record_suppressed_indices:
            return dets, suppressed_indices
        else:
            return dets


def bbox_decode(priors, bbox_preds):
    xys = (bbox_preds[..., :2] * priors[:, 2:]) + priors[:, :2]
    whs = bbox_preds[..., 2:].exp() * priors[:, 2:]

    tl_x = (xys[..., 0] - whs[..., 0] / 2)
    tl_y = (xys[..., 1] - whs[..., 1] / 2)
    br_x = (xys[..., 0] + whs[..., 0] / 2)
    br_y = (xys[..., 1] + whs[..., 1] / 2)

    decoded_bboxes = torch.stack([tl_x, tl_y, br_x, br_y], -1)
    return decoded_bboxes