from __future__ import annotations

import torch
from torch import Tensor


def box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor:
    cx, cy, width, height = boxes.unbind(-1)
    return torch.stack(
        (cx - 0.5 * width, cy - 0.5 * height, cx + 0.5 * width, cy + 0.5 * height),
        dim=-1,
    )


def box_xyxy_to_cxcywh(boxes: Tensor) -> Tensor:
    x0, y0, x1, y1 = boxes.unbind(-1)
    return torch.stack(
        ((x0 + x1) * 0.5, (y0 + y1) * 0.5, x1 - x0, y1 - y0), dim=-1
    )


def box_area(boxes: Tensor) -> Tensor:
    return (boxes[..., 2] - boxes[..., 0]).clamp(min=0) * (
        boxes[..., 3] - boxes[..., 1]
    ).clamp(min=0)


def box_iou(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]:
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)
    top_left = torch.maximum(boxes1[:, None, :2], boxes2[:, :2])
    bottom_right = torch.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
    intersection = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
    union = area1[:, None] + area2 - intersection
    return intersection / union.clamp(min=1e-7), union


def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """Pairwise generalized IoU for boxes in x0, y0, x1, y1 format."""
    iou, union = box_iou(boxes1, boxes2)
    top_left = torch.minimum(boxes1[:, None, :2], boxes2[:, :2])
    bottom_right = torch.maximum(boxes1[:, None, 2:], boxes2[:, 2:])
    enclosing = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
    return iou - (enclosing - union) / enclosing.clamp(min=1e-7)


def generalized_box_iou_batched(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """Batched pairwise generalized IoU, boxes in x0, y0, x1, y1 format.

    boxes1: [B, N, 4], boxes2: [B, M, 4] -> [B, N, M]
    """
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)
    top_left = torch.maximum(boxes1[:, :, None, :2], boxes2[:, None, :, :2])
    bottom_right = torch.minimum(boxes1[:, :, None, 2:], boxes2[:, None, :, 2:])
    intersection = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
    union = area1[:, :, None] + area2[:, None, :] - intersection
    iou = intersection / union.clamp(min=1e-7)
    enc_top_left = torch.minimum(boxes1[:, :, None, :2], boxes2[:, None, :, :2])
    enc_bottom_right = torch.maximum(boxes1[:, :, None, 2:], boxes2[:, None, :, 2:])
    enclosing = (enc_bottom_right - enc_top_left).clamp(min=0).prod(dim=-1)
    return iou - (enclosing - union) / enclosing.clamp(min=1e-7)


def inverse_sigmoid(value: Tensor, eps: float = 1e-5) -> Tensor:
    value = value.clamp(min=0.0, max=1.0)
    return torch.log(value.clamp(min=eps) / (1.0 - value).clamp(min=eps))