"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""

from enum import Enum

import torch
import torchvision
from torch import Tensor


class BoxProcessFormat(Enum):
    """Box process format

    Available formats are
    * ``RESIZE``
    * ``RESIZE_KEEP_RATIO``
    * ``RESIZE_KEEP_RATIO_PADDING``
    """

    RESIZE = 1
    RESIZE_KEEP_RATIO = 2
    RESIZE_KEEP_RATIO_PADDING = 3


def box_revert(
    boxes: Tensor,
    orig_sizes: Tensor = None,
    eval_sizes: Tensor = None,
    inpt_sizes: Tensor = None,
    inpt_padding: Tensor = None,
    normalized: bool = True,
    in_fmt: str = "cxcywh",
    out_fmt: str = "xyxy",
    process_fmt=BoxProcessFormat.RESIZE,
) -> Tensor:
    """
    Args:
        boxes(Tensor), [N, :, 4], (x1, y1, x2, y2), pred boxes.
        inpt_sizes(Tensor), [N, 2], (w, h). input sizes.
        orig_sizes(Tensor), [N, 2], (w, h). origin sizes.
        inpt_padding (Tensor), [N, 2], (w_pad, h_pad, ...).
        (inpt_sizes + inpt_padding) == eval_sizes
    """
    assert in_fmt in ("cxcywh", "xyxy"), ""

    if normalized and eval_sizes is not None:
        boxes = boxes * eval_sizes.repeat(1, 2).unsqueeze(1)

    if inpt_padding is not None:
        if in_fmt == "xyxy":
            boxes -= inpt_padding[:, :2].repeat(1, 2).unsqueeze(1)
        elif in_fmt == "cxcywh":
            boxes[..., :2] -= inpt_padding[:, :2].repeat(1, 2).unsqueeze(1)

    if orig_sizes is not None:
        orig_sizes = orig_sizes.repeat(1, 2).unsqueeze(1)
        if inpt_sizes is not None:
            inpt_sizes = inpt_sizes.repeat(1, 2).unsqueeze(1)
            boxes = boxes * (orig_sizes / inpt_sizes)
        else:
            boxes = boxes * orig_sizes

    boxes = torchvision.ops.box_convert(boxes, in_fmt=in_fmt, out_fmt=out_fmt)
    return boxes