from __future__ import annotations

from typing import Iterable, List, Sequence, Tuple

import cv2
import numpy as np


# DOTAv1 classes, in the order encoded by the exported model.
CLASS_NAMES: Tuple[str, ...] = (
    "plane", "ship", "storage tank", "baseball diamond", "tennis court",
    "basketball court", "ground track field", "harbor", "bridge",
    "large vehicle", "small vehicle", "helicopter", "roundabout",
    "soccer ball field", "swimming pool",
)

# One BGR colour per class; distinct enough to tell adjacent classes apart.
CLASS_COLORS: Tuple[Tuple[int, int, int], ...] = (
    (0, 215, 255), (0, 140, 255), (255, 200, 0), (255, 120, 200),
    (90, 230, 0), (255, 90, 140), (200, 200, 0), (128, 0, 255),
    (0, 255, 180), (255, 0, 0), (0, 0, 255), (60, 76, 231),
    (255, 255, 0), (0, 165, 255), (147, 20, 255),
)

LABEL_FG = (255, 255, 255)
FONT = cv2.FONT_HERSHEY_SIMPLEX


def letterbox(
    image: np.ndarray,
    input_size: int = 1024,
    pad_value: int = 114,
) -> Tuple[np.ndarray, float, int, int]:
    """Ultralytics-style letterbox: keep ratio, centered gray padding."""
    height, width = image.shape[:2]
    ratio = min(input_size / width, input_size / height)
    resized_w, resized_h = round(width * ratio), round(height * ratio)
    resized = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR)
    pad_w, pad_h = input_size - resized_w, input_size - resized_h
    left, top = pad_w // 2, pad_h // 2
    canvas = cv2.copyMakeBorder(
        resized,
        top,
        pad_h - top,
        left,
        pad_w - left,
        cv2.BORDER_CONSTANT,
        value=(pad_value, pad_value, pad_value),
    )
    return canvas, ratio, left, top


def preprocess_image(
    image: np.ndarray,
    input_size: int = 1024,
    pad_value: int = 114,
) -> Tuple[np.ndarray, float, int, int]:
    canvas, ratio, pad_x, pad_y = letterbox(image, input_size, pad_value)
    rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
    tensor = (rgb.astype(np.float32) / 255.0)[None, :]
    return tensor, ratio, pad_x, pad_y


def xywhr2xyxyxyxy(boxes: np.ndarray) -> np.ndarray:
    """Convert [cx, cy, w, h, angle] boxes to (N, 4, 2) corner points."""
    boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 5)
    center = boxes[:, :2]
    w, h, angle = boxes[:, 2:3], boxes[:, 3:4], boxes[:, 4:5]
    cos, sin = np.cos(angle), np.sin(angle)
    vec1 = np.concatenate([w / 2 * cos, w / 2 * sin], axis=-1)
    vec2 = np.concatenate([-h / 2 * sin, h / 2 * cos], axis=-1)
    pt1 = center + vec1 + vec2
    pt2 = center + vec1 - vec2
    pt3 = center - vec1 - vec2
    pt4 = center - vec1 + vec2
    return np.stack([pt1, pt2, pt3, pt4], axis=1)


def _covariance_matrix(boxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Gaussian covariance components (a, b, c) for xywhr boxes."""
    w2 = boxes[:, 2] ** 2 / 12.0
    h2 = boxes[:, 3] ** 2 / 12.0
    angle = boxes[:, 4]
    cos2 = np.cos(angle) ** 2
    sin2 = np.sin(angle) ** 2
    a = w2 * cos2 + h2 * sin2
    b = w2 * sin2 + h2 * cos2
    c = (w2 - h2) * np.cos(angle) * np.sin(angle)
    return a, b, c


def batch_probiou(obb1: np.ndarray, obb2: np.ndarray, eps: float = 1e-7) -> np.ndarray:
    """Probabilistic IoU between two sets of xywhr boxes -> shape (N, M).

    Mirrors ``ultralytics.utils.metrics.batch_probiou`` so on-device NMS
    matches the official OBB post-processing.
    """
    obb1 = np.asarray(obb1, dtype=np.float32).reshape(-1, 5)
    obb2 = np.asarray(obb2, dtype=np.float32).reshape(-1, 5)
    x1, y1 = obb1[:, 0:1], obb1[:, 1:2]
    x2, y2 = obb2[None, :, 0], obb2[None, :, 1]
    a1, b1, c1 = (v[:, None] for v in _covariance_matrix(obb1))
    a2, b2, c2 = (v[None, :] for v in _covariance_matrix(obb2))

    denom = (a1 + a2) * (b1 + b2) - (c1 + c2) ** 2 + eps
    t1 = (((a1 + a2) * (y1 - y2) ** 2 + (b1 + b2) * (x1 - x2) ** 2) / denom) * 0.25
    t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / denom) * 0.5
    det = np.sqrt(
        np.clip(a1 * b1 - c1 ** 2, 0, None) * np.clip(a2 * b2 - c2 ** 2, 0, None)
    )
    t3 = np.log(((a1 + a2) * (b1 + b2) - (c1 + c2) ** 2) / (4 * det + eps) + eps) * 0.5
    bd = np.clip(t1 + t2 + t3, eps, 100.0)
    hd = np.sqrt(1.0 - np.exp(-bd) + eps)
    return 1.0 - hd


def obb_nms(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> List[int]:
    """Greedy NMS over xywhr boxes using probabilistic IoU."""
    if len(boxes) == 0:
        return []
    order = scores.argsort()[::-1]
    keep: List[int] = []
    while order.size:
        i = int(order[0])
        keep.append(i)
        if order.size == 1:
            break
        rest = order[1:]
        iou = batch_probiou(boxes[i:i + 1], boxes[rest])[0]
        order = rest[iou <= iou_thres]
    return keep


def obb_postprocess(
    prediction: np.ndarray,
    ratio: float,
    pad_x: int,
    pad_y: int,
    original_shape: Sequence[int],
    conf_thres: float = 0.25,
    iou_thres: float = 0.45,
    max_det: int = 300,
    class_names: Sequence[str] = CLASS_NAMES,
    class_agnostic: bool = False,
) -> List[dict]:
    """Decode output0 `[1, 4+nc+1, anchors]` into oriented boxes.

    Channel order matches the Ultralytics OBB head: xywh, class scores, angle.
    """
    prediction = np.asarray(prediction, dtype=np.float32)
    if prediction.ndim == 3:
        if prediction.shape[0] != 1:
            raise ValueError(f"Only batch size 1 is supported, got {prediction.shape}")
        prediction = prediction[0]
    if prediction.shape[0] < prediction.shape[1]:
        prediction = prediction.T  # channel-first -> (anchors, channels)
    if prediction.ndim != 2 or prediction.shape[1] < 6:
        raise ValueError(f"Unexpected prediction shape: {prediction.shape}")

    nc = prediction.shape[1] - 5
    class_scores = prediction[:, 4:4 + nc]
    class_ids = class_scores.argmax(axis=1)
    scores = class_scores[np.arange(class_scores.shape[0]), class_ids]
    mask = scores >= conf_thres
    if not np.any(mask):
        return []

    xywh = prediction[mask, :4]
    angle = prediction[mask, 4 + nc]
    scores = scores[mask]
    class_ids = class_ids[mask]
    rboxes = np.concatenate([xywh, angle[:, None]], axis=1)

    # Class offsets keep NMS from suppressing different classes, like Ultralytics.
    nms_boxes = rboxes.copy()
    if not class_agnostic:
        offsets = class_ids.astype(np.float32) * 7680.0
        nms_boxes[:, 0] += offsets
        nms_boxes[:, 1] += offsets
    keep = obb_nms(nms_boxes, scores, iou_thres)[:max_det]

    boxes = rboxes[keep].copy()
    boxes[:, 0] = (boxes[:, 0] - pad_x) / ratio
    boxes[:, 1] = (boxes[:, 1] - pad_y) / ratio
    boxes[:, 2] = boxes[:, 2] / ratio
    boxes[:, 3] = boxes[:, 3] / ratio
    corners = xywhr2xyxyxyxy(boxes)
    corners[..., 0].clip(0, original_shape[1], out=corners[..., 0])
    corners[..., 1].clip(0, original_shape[0], out=corners[..., 1])

    detections: List[dict] = []
    for row, idx in enumerate(keep):
        cls_id = int(class_ids[idx])
        detections.append({
            "box": boxes[row].astype(np.float32),
            "corners": corners[row].astype(np.float32),
            "score": float(scores[idx]),
            "class_id": cls_id,
            "class_name": class_names[cls_id] if cls_id < len(class_names) else str(cls_id),
        })
    return detections


def draw_label(result: np.ndarray, text: str, x: int, bottom: int, color) -> None:
    """Filled label above ``(x, bottom)``; drops inside the image when it clips."""
    scale, thickness, pad = 0.5, 1, 4
    (text_w, text_h), _ = cv2.getTextSize(text, FONT, scale, thickness)
    height = text_h + 2 * pad
    top = bottom - height
    if top < 0:  # no room above the box: draw the label just inside its top edge
        top = bottom
    right = min(result.shape[1] - 1, x + text_w + 2 * pad)
    cv2.rectangle(result, (x, top), (right, top + height), color, -1, cv2.LINE_AA)
    cv2.putText(result, text, (x + pad, top + text_h + pad), FONT, scale,
                LABEL_FG, thickness, cv2.LINE_AA)


def draw_obb_res(image: np.ndarray, detections: Iterable[dict]) -> np.ndarray:
    result = image.astype(np.uint8).copy()
    for i, det in enumerate(detections, start=1):
        corners = np.asarray(det["corners"], dtype=np.float32)
        score = float(det["score"])
        name = det["class_name"]
        color = CLASS_COLORS[det["class_id"] % len(CLASS_COLORS)]
        points = np.round(corners).astype(np.int32).reshape(-1, 1, 2)
        cv2.polylines(result, [points], True, color, 2, cv2.LINE_AA)
        top_x, top_y = np.round(corners[np.argmin(corners[:, 1])]).astype(int)
        draw_label(result, f"{name} {score:.2f}", int(top_x), int(top_y), color)
        print(i, det["class_id"], name, [round(float(v), 2) for v in det["box"]], score)
    return result