from __future__ import annotations

import numpy as np
import numpy.typing as npt
from deprecate import TargetMode, deprecated

from supervision.detection.utils.iou_and_nms import box_iou_batch


def clip_boxes(
    xyxy: npt.NDArray[np.number],
    resolution_wh: tuple[int, int],
) -> npt.NDArray[np.number]:
    """
    Clips bounding boxes coordinates to fit within the frame resolution.

    Args:
        xyxy: A numpy array of shape `(N, 4)` where each
            row corresponds to a bounding box in
            the format `(x_min, y_min, x_max, y_max)`.
        resolution_wh: A tuple of the form
            `(width, height)` representing the resolution of the frame.

    Returns:
        A numpy array of shape `(N, 4)` where each row
            corresponds to a bounding box with coordinates clipped to fit
            within the frame resolution.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> import supervision as sv
        >>> xyxy = np.array([
        ...     [10, 20, 300, 200],
        ...     [15, 25, 350, 450],
        ...     [-10, -20, 30, 40]
        ... ])
        >>> sv.clip_boxes(xyxy=xyxy, resolution_wh=(320, 240))
        array([[ 10,  20, 300, 200],
               [ 15,  25, 320, 240],
               [  0,   0,  30,  40]])

        ```
    """
    result: npt.NDArray[np.number] = np.copy(xyxy)
    width, height = resolution_wh
    result[:, [0, 2]] = result[:, [0, 2]].clip(0, width)
    result[:, [1, 3]] = result[:, [1, 3]].clip(0, height)
    return result


def pad_boxes(
    xyxy: npt.NDArray[np.number],
    px: int,
    py: int | None = None,
) -> npt.NDArray[np.number]:
    """
    Pads bounding boxes coordinates with a constant padding.

    Args:
        xyxy: A numpy array of shape `(N, 4)` where each
            row corresponds to a bounding box in the format
            `(x_min, y_min, x_max, y_max)`.
        px: The padding value to be added to both the left and right sides of
            each bounding box.
        py: The padding value to be added to both the top and bottom
            sides of each bounding box. If not provided, `px` will be used for both
            dimensions.

    Returns:
        A numpy array of shape `(N, 4)` where each row corresponds to a
            bounding box with coordinates padded according to the provided padding
            values.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> import supervision as sv
        >>> xyxy = np.array([
        ...     [10, 20, 30, 40],
        ...     [15, 25, 35, 45]
        ... ])
        >>> sv.pad_boxes(xyxy=xyxy, px=5, py=10)
        array([[ 5, 10, 35, 50],
               [10, 15, 40, 55]])

        ```
    """
    if py is None:
        py = px

    result = xyxy.copy()
    result[:, [0, 1]] -= [px, py]
    result[:, [2, 3]] += [px, py]

    return result


@deprecated(  # type: ignore[untyped-decorator]
    target=TargetMode.ARGS_REMAP,
    deprecated_in="0.27.0",
    remove_in="0.30.0",
    args_mapping={"normalized_xyxy": "xyxy"},
)
def denormalize_boxes(
    xyxy: npt.NDArray[np.number],
    resolution_wh: tuple[int, int],
    normalization_factor: float = 1.0,
    normalized_xyxy: npt.NDArray[np.number] | None = None,
) -> npt.NDArray[np.number]:
    """
    Convert normalized bounding box coordinates to absolute pixel coordinates.

    Multiplies each bounding box coordinate by image size and divides by
    `normalization_factor`, mapping values from normalized `[0, normalization_factor]`
    to absolute pixel values for a given resolution.

    Args:
        xyxy: Normalized bounding boxes of shape `(N, 4)`,
            where each row is `(x_min, y_min, x_max, y_max)`, values in
            `[0, normalization_factor]`.
        resolution_wh: Target image resolution as `(width, height)`.
        normalization_factor: Maximum value of input coordinate range.
            Defaults to `1.0`.

    Returns:
        Array of shape `(N, 4)` with absolute coordinates in
            `(x_min, y_min, x_max, y_max)` format.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> import supervision as sv
        >>> xyxy = np.array([
        ...     [0.1, 0.2, 0.5, 0.6],
        ...     [0.3, 0.4, 0.7, 0.8],
        ...     [0.2, 0.1, 0.6, 0.5]
        ... ])
        >>> sv.denormalize_boxes(xyxy, (1280, 720))
        array([[128., 144., 640., 432.],
               [384., 288., 896., 576.],
               [256.,  72., 768., 360.]])

        ```

        ```pycon
        >>> xyxy = np.array([
        ...     [256., 128., 768., 640.]
        ... ])
        >>> sv.denormalize_boxes(xyxy, (1280, 720), normalization_factor=1024.0)
        array([[320.,  90., 960., 450.]])

        ```
    """
    width, height = resolution_wh
    result = xyxy.copy()

    result[:, [0, 2]] = (result[:, [0, 2]] * width) / normalization_factor
    result[:, [1, 3]] = (result[:, [1, 3]] * height) / normalization_factor

    return result


def move_boxes(
    xyxy: npt.NDArray[np.float64], offset: npt.NDArray[np.int32]
) -> npt.NDArray[np.float64]:
    """
    Args:
        xyxy: An array of shape `(n, 4)` containing the
            bounding boxes coordinates in format `[x1, y1, x2, y2]`
        offset: An array of shape `(2,)` containing offset values in format
            is `[dx, dy]`.

    Returns:
        Repositioned bounding boxes.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> import supervision as sv
        >>> xyxy = np.array([
        ...     [10, 10, 20, 20],
        ...     [30, 30, 40, 40]
        ... ])
        >>> offset = np.array([5, 5])
        >>> sv.move_boxes(xyxy=xyxy, offset=offset)
        array([[15, 15, 25, 25],
               [35, 35, 45, 45]])

        ```
    """
    return xyxy + np.hstack([offset, offset])


def move_oriented_boxes(
    xyxyxyxy: npt.NDArray[np.float64], offset: npt.NDArray[np.int32]
) -> npt.NDArray[np.float64]:
    """
    Args:
        xyxyxyxy: An array of shape `(n, 4, 2)` containing the
        oriented bounding boxes coordinates in format
        `[[x1, y1], [x2, y2], [x3, y3], [x3, y3]]`
        offset: An array of shape `(2,)` containing offset values in format
            is `[dx, dy]`.

    Returns:
        Repositioned bounding boxes.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> from supervision.detection.utils.boxes import move_oriented_boxes
        >>> xyxyxyxy = np.array([
        ...     [
        ...         [20, 10],
        ...         [10, 20],
        ...         [20, 30],
        ...         [30, 20]
        ...     ],
        ...     [
        ...         [30, 30],
        ...         [20, 40],
        ...         [30, 50],
        ...         [40, 40]
        ...     ]
        ... ])
        >>> offset = np.array([5, 5])
        >>> move_oriented_boxes(xyxyxyxy=xyxyxyxy, offset=offset)
        array([[[25, 15],
                [15, 25],
                [25, 35],
                [35, 25]],
        <BLANKLINE>
               [[35, 35],
                [25, 45],
                [35, 55],
                [45, 45]]])

        ```
    """
    return xyxyxyxy + offset


def obb_polygon_area(corners: npt.NDArray) -> npt.NDArray[np.float64]:
    """Compute the area of N oriented bounding boxes using the shoelace formula.

    Args:
        corners: OBB corner coordinates with shape `(N, 4, 2)`.

    Returns:
        Area of each box as a 1-D float64 array of shape `(N,)`.

    Raises:
        ValueError: If `corners` does not have shape `(N, 4, 2)`.

    Examples:
        >>> import numpy as np
        >>> from supervision.detection.utils.boxes import obb_polygon_area
        >>> corners = np.array([[[0, 5], [5, 10], [10, 5], [5, 0]]], dtype=np.float32)
        >>> obb_polygon_area(corners)
        array([50.])
    """
    corners = np.asarray(corners)
    if corners.ndim != 3 or corners.shape[-2:] != (4, 2):
        raise ValueError(f"corners must have shape (N, 4, 2); got {corners.shape}")
    x = corners[..., 0].astype(np.float64, copy=False)
    y = corners[..., 1].astype(np.float64, copy=False)
    cross = x * np.roll(y, -1, axis=-1) - y * np.roll(x, -1, axis=-1)
    return 0.5 * np.abs(np.sum(cross, axis=-1))


def xyxyxyxy_to_xyxy(
    xyxyxyxy: npt.NDArray[np.number],
) -> npt.NDArray[np.number]:
    """Convert oriented bounding box corners to axis-aligned bounding boxes.

    Args:
        xyxyxyxy: OBB corner coordinates with shape `(N, 4, 2)` where each
            box is represented as `[[x1, y1], [x2, y2], [x3, y3], [x4, y4]]`.

    Returns:
        Axis-aligned bounding boxes as an array of shape `(N, 4)`
            in `(x_min, y_min, x_max, y_max)` format.

    Raises:
        ValueError: If `xyxyxyxy` does not have shape `(N, 4, 2)`.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> import supervision as sv
        >>> corners = np.array([
        ...     [[0, 0], [10, 0], [10, 5], [0, 5]],
        ...     [[5, 5], [15, 5], [15, 10], [5, 10]],
        ... ], dtype=np.float32)
        >>> sv.xyxyxyxy_to_xyxy(corners)
        array([[ 0.,  0., 10.,  5.],
               [ 5.,  5., 15., 10.]], dtype=float32)

        ```
    """
    xyxyxyxy = np.asarray(xyxyxyxy)
    if xyxyxyxy.ndim != 3 or xyxyxyxy.shape[-2:] != (4, 2):
        raise ValueError(f"xyxyxyxy must have shape (N, 4, 2); got {xyxyxyxy.shape}")
    x_min = xyxyxyxy[..., 0].min(axis=-1)
    y_min = xyxyxyxy[..., 1].min(axis=-1)
    x_max = xyxyxyxy[..., 0].max(axis=-1)
    y_max = xyxyxyxy[..., 1].max(axis=-1)
    return np.stack([x_min, y_min, x_max, y_max], axis=-1)


def scale_boxes(
    xyxy: npt.NDArray[np.float64], factor: float
) -> npt.NDArray[np.float64]:
    """
    Scale the dimensions of bounding boxes.

    Args:
        xyxy: An array of shape `(n, 4)` containing the
            bounding boxes coordinates in format `[x1, y1, x2, y2]`
        factor: A float value representing the factor by which the box
            dimensions are scaled. A factor greater than 1 enlarges the boxes, while a
            factor less than 1 shrinks them.

    Returns:
        Scaled bounding boxes.

    Examples:
        ```pycon
        >>> import numpy as np
        >>> import supervision as sv
        >>> xyxy = np.array([
        ...     [10, 10, 20, 20],
        ...     [30, 30, 40, 40]
        ... ])
        >>> sv.scale_boxes(xyxy=xyxy, factor=1.5)
        array([[ 7.5,  7.5, 22.5, 22.5],
               [27.5, 27.5, 42.5, 42.5]])

        ```
    """
    centers = (xyxy[:, :2] + xyxy[:, 2:]) / 2
    new_sizes = (xyxy[:, 2:] - xyxy[:, :2]) * factor
    return np.concatenate((centers - new_sizes / 2, centers + new_sizes / 2), axis=1)


def spread_out_boxes(
    xyxy: npt.NDArray[np.number],
    max_iterations: int = 100,
) -> npt.NDArray[np.number]:
    """
    Spread out boxes that overlap with each other.

    Args:
        xyxy: Numpy array of shape (N, 4) where N is the number of boxes.
        max_iterations: Maximum number of iterations to run the algorithm for.

    Example:
        ```pycon
        >>> import numpy as np
        >>> from supervision.detection.utils.boxes import spread_out_boxes
        >>> xyxy = np.array([
        ...     [10, 10, 20, 20],
        ...     [12, 12, 22, 22]
        ... ])
        >>> spread_out = spread_out_boxes(xyxy=xyxy, max_iterations=10)
        >>> # The boxes should be moved apart
        >>> bool(spread_out[0, 0] < 10 and spread_out[0, 1] < 10)
        True
        >>> bool(spread_out[1, 0] > 12 and spread_out[1, 1] > 12)
        True

        ```
    """
    if len(xyxy) == 0:
        return xyxy

    xyxy_padded = pad_boxes(xyxy, px=1)
    for _ in range(max_iterations):
        # NxN
        iou = box_iou_batch(xyxy_padded, xyxy_padded)
        np.fill_diagonal(iou, 0)
        if np.all(iou == 0):
            break

        overlap_mask = iou > 0

        # Nx2
        centers = (xyxy_padded[:, :2] + xyxy_padded[:, 2:]) / 2

        # NxNx2
        delta_centers = centers[:, np.newaxis, :] - centers[np.newaxis, :, :]
        delta_centers *= overlap_mask[:, :, np.newaxis]

        # Nx2
        delta_sum = np.sum(delta_centers, axis=1)
        delta_magnitude = np.linalg.norm(delta_sum, axis=1, keepdims=True)
        direction_vectors = np.divide(
            delta_sum,
            delta_magnitude,
            out=np.zeros_like(delta_sum),
            where=delta_magnitude != 0,
        )

        force_vectors = np.sum(iou, axis=1)
        force_vectors = force_vectors[:, np.newaxis] * direction_vectors

        force_vectors *= 10
        force_vectors[(force_vectors > 0) & (force_vectors < 2)] = 2
        force_vectors[(force_vectors < 0) & (force_vectors > -2)] = -2

        force_vectors = force_vectors.astype(int)

        xyxy_padded[:, [0, 1]] += force_vectors
        xyxy_padded[:, [2, 3]] += force_vectors

    return pad_boxes(xyxy_padded, px=-1)