#!/usr/bin/env python3
"""
YOLO11-OBB axmodel inference (AXEngine NPU).

Usage:
    python3 ax_infer.py
    python3 ax_infer.py -m yolo11n-obb_1024x1024.axmodel -i boats.jpg -o result_ax.jpg
"""

import argparse
import math
import os
from time import time

import cv2
import numpy as np
import axengine as axe


DOTA_CLASSES = [
    "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",
]

DOTA_COLORS = [
    (255, 56, 56),   (255, 159, 56),  (255, 207, 56),  (180, 255, 56),
    (102, 255, 56),  (56, 255, 122),  (56, 255, 207),  (56, 207, 255),
    (56, 122, 255),  (102, 56, 255),  (180, 56, 255),  (255, 56, 207),
    (255, 56, 122),  (200, 200, 200), (128, 128, 255),
]


def infer_hw_layout(shape):
    s = list(shape)
    if len(s) == 4 and s[-1] == 3:
        return int(s[1] or 1024), int(s[2] or 1024), "NHWC"
    if len(s) == 4 and s[1] == 3:
        return int(s[2] or 1024), int(s[3] or 1024), "NCHW"
    return 1024, 1024, "NCHW"


def letterbox(image, new_shape, pad=114):
    h0, w0 = image.shape[:2]
    H, W = new_shape
    r = min(H / h0, W / w0)
    nw, nh = round(w0 * r), round(h0 * r)
    dw, dh = (W - nw) / 2.0, (H - nh) / 2.0
    if (w0, h0) != (nw, nh):
        image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_LINEAR)
    top, bottom = round(dh - 0.1), round(dh + 0.1)
    left, right = round(dw - 0.1), round(dw + 0.1)
    image = cv2.copyMakeBorder(image, top, bottom, left, right,
                               cv2.BORDER_CONSTANT, value=(pad,) * 3)
    return image, r, (left, top)


def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)


def dfl_decode(box, reg_max):
    box = box.reshape(-1, 4, reg_max)
    return np.sum(softmax(box, axis=-1) * np.arange(reg_max, dtype=np.float32), axis=-1)


def decode_rbox(box, angle, anchors, stride, reg_max):
    if reg_max > 1:
        box = dfl_decode(box, reg_max)
    cos_a, sin_a = np.cos(angle), np.sin(angle)
    xf = (box[:, 2] - box[:, 0]) * 0.5
    yf = (box[:, 3] - box[:, 1]) * 0.5
    cx = xf * cos_a - yf * sin_a + anchors[:, 0]
    cy = xf * sin_a + yf * cos_a + anchors[:, 1]
    w = box[:, 0] + box[:, 2]
    h = box[:, 1] + box[:, 3]
    return np.stack([cx * stride, cy * stride, w * stride, h * stride, angle], axis=1)


def _covariance(boxes):
    a = boxes[:, 2] ** 2 / 12.0
    b = boxes[:, 3] ** 2 / 12.0
    c = boxes[:, 4]
    cos2, sin2 = np.cos(c) ** 2, np.sin(c) ** 2
    return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * np.cos(c) * np.sin(c)


def probiou(obb1, obb2, eps=1e-7):
    x1, y1 = obb1[:, 0:1], obb1[:, 1:2]
    x2, y2 = obb2[:, 0][None, :], obb2[:, 1][None, :]
    a1, b1, c1 = (v[:, None] for v in _covariance(obb1))
    a2f, b2f, c2f = _covariance(obb2)
    a2, b2, c2 = a2f[None, :], b2f[None, :], c2f[None, :]
    s = (a1 + a2) * (b1 + b2) - (c1 + c2) ** 2
    t1 = ((a1 + a2) * (y1 - y2) ** 2 + (b1 + b2) * (x1 - x2) ** 2) / (s + eps) * 0.25
    t2 = ((c1 + c2) * (x2 - x1) * (y1 - y2)) / (s + eps) * 0.5
    inner = ((a1 * b1 - c1 ** 2).clip(0) * (a2 * b2 - c2 ** 2).clip(0))
    t3 = np.log(s / (4.0 * np.sqrt(inner) + eps) + eps) * 0.5
    bd = np.clip(t1 + t2 + t3, eps, 100.0)
    return 1.0 - np.sqrt(1.0 - np.exp(-bd) + eps)


def nms_rotated(rboxes, scores, classes, iou_thres, max_wh=7680.0, agnostic=False):
    if rboxes.size == 0:
        return np.empty((0,), dtype=np.int64)
    b = rboxes.copy()
    if not agnostic:
        off = classes.astype(np.float32) * max_wh
        b[:, 0] += off
        b[:, 1] += off
    order = np.argsort(-scores)
    ious = probiou(b[order], b[order])
    n = ious.shape[0]
    ious *= np.triu(np.ones((n, n), dtype=bool), k=1)
    keep = (ious >= iou_thres).sum(axis=0) <= 0
    return order[keep]


def rbox_corners(cx, cy, w, h, ag):
    c, s = math.cos(ag), math.sin(ag)
    wx, wy = w / 2 * c, w / 2 * s
    hx, hy = -h / 2 * s, h / 2 * c
    return np.array([
        [cx - wx - hx, cy - wy - hy],
        [cx + wx - hx, cy + wy - hy],
        [cx + wx + hx, cy + wy + hy],
        [cx - wx + hx, cy - wy + hy],
    ], dtype=np.float32)


def sort_outputs(names):
    order = {"box": 0, "cls": 1, "ang": 2}
    return sorted(names, key=lambda n: (int(n.split("_")[1][5:]), order[n.split("_")[0]]))


def postprocess(outs, names, score_thres, nms_thres, agnostic):
    conf_raw = -math.log(1.0 / score_thres - 1.0)
    rboxes, scores, classes = [], [], []

    for i in range(len(outs) // 3):
        box, cls, ang = outs[i * 3], outs[i * 3 + 1], outs[i * 3 + 2]
        h, w = box.shape[1:3]
        stride = int(names[i * 3].split("stride")[-1])

        box = box[0].reshape(-1, box.shape[-1])
        cls = cls[0].reshape(-1, cls.shape[-1])
        ang = ang[0].reshape(-1, ang.shape[-1])

        reg_max = box.shape[-1] // 4

        cls_max = cls.max(axis=1) if cls.shape[-1] > 1 else cls[:, 0]
        cls_id = cls.argmax(axis=1) if cls.shape[-1] > 1 else np.zeros(len(cls_max), dtype=np.int32)

        keep = cls_max >= conf_raw
        if not keep.any():
            continue

        gy, gx = np.indices((h, w))
        anchors = (np.stack((gx.ravel(), gy.ravel()), -1).astype(np.float32) + 0.5)[keep]

        # YOLO11-OBB angle: (sigmoid(raw) - 0.25) * pi
        angle = (1.0 / (1.0 + np.exp(-ang[keep])) - 0.25) * np.pi
        angle = angle.reshape(-1)

        rboxes.append(decode_rbox(box[keep], angle, anchors, stride, reg_max))
        scores.append(1.0 / (1.0 + np.exp(-cls_max[keep])))
        classes.append(cls_id[keep])

    if not rboxes:
        return np.empty((0, 5), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32)

    rboxes = np.concatenate(rboxes).astype(np.float32)
    scores = np.concatenate(scores).astype(np.float32)
    classes = np.concatenate(classes).astype(np.int32)
    k = nms_rotated(rboxes, scores, classes, nms_thres, agnostic=agnostic)
    return rboxes[k], scores[k], classes[k]


def scale_back(rboxes, gain, pad, orig_shape):
    rboxes = rboxes.copy()
    rboxes[:, 0] -= pad[0]
    rboxes[:, 1] -= pad[1]
    rboxes[:, :4] /= gain
    rboxes[:, 0] = np.clip(rboxes[:, 0], 0, orig_shape[1])
    rboxes[:, 1] = np.clip(rboxes[:, 1], 0, orig_shape[0])
    t = np.mod(rboxes[:, 4], np.pi)
    swap = t >= (np.pi / 2)
    if swap.any():
        rboxes[swap, 2], rboxes[swap, 3] = rboxes[swap, 3].copy(), rboxes[swap, 2].copy()
    rboxes[:, 4] = np.mod(rboxes[:, 4], np.pi / 2)
    return rboxes


def draw(img, rboxes, scores, classes):
    for (cx, cy, w, h, th), sc, cid in zip(rboxes, scores, classes):
        name = DOTA_CLASSES[cid] if cid < len(DOTA_CLASSES) else f"cls{cid}"
        color = DOTA_COLORS[cid % len(DOTA_COLORS)]
        pts = rbox_corners(cx, cy, w, h, th).astype(np.int32)
        cv2.polylines(img, [pts], True, color, 2, cv2.LINE_AA)
        label = f"{name} {sc:.2f}"
        (tw, tht), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        x, y = int(pts[0][0]), max(0, int(pts[0][1]) - 5)
        cv2.rectangle(img, (x, y - tht - 2), (x + tw + 2, y + 2), color, -1)
        cv2.putText(img, label, (x + 1, y - 1), cv2.FONT_HERSHEY_SIMPLEX,
                    0.5, (255, 255, 255), 1, cv2.LINE_AA)


def main():
    ap = argparse.ArgumentParser(description="YOLO11-OBB axmodel inference (AXEngine NPU)")
    ap.add_argument("-m", "--model", default="yolo11n-obb_1024x1024.axmodel")
    ap.add_argument("-i", "--img", default="boats.jpg")
    ap.add_argument("-o", "--output", default="result_yolo11_obb_ax.jpg")
    ap.add_argument("--score-thres", type=float, default=0.25)
    ap.add_argument("--nms-thres", type=float, default=0.45)
    ap.add_argument("--max-det", type=int, default=300)
    ap.add_argument("--agnostic-nms", action="store_true")
    ap.add_argument("--providers", default="AxEngineExecutionProvider")
    opt = ap.parse_args()

    assert os.path.exists(opt.model), f"model missing: {opt.model}"
    assert os.path.exists(opt.img), f"image missing: {opt.img}"

    t0 = time()
    providers = [p.strip() for p in opt.providers.split(",") if p.strip()] or None
    sess = axe.InferenceSession(opt.model, providers=providers)
    print(f"Load model: {(time()-t0)*1000:.1f} ms")

    inp = sess.get_inputs()[0]
    H, W, layout = infer_hw_layout(inp.shape)

    img0 = cv2.imread(opt.img)
    assert img0 is not None, f"cannot read {opt.img}"

    img, r, pad = letterbox(img0.copy(), (H, W))
    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    tensor = (rgb[None, ...].astype(np.uint8) if layout == "NHWC"
              else np.transpose(rgb, (2, 0, 1))[None, ...].astype(np.uint8))

    t0 = time()
    names = sort_outputs([o.name for o in sess.get_outputs()])
    outs = sess.run(names, {inp.name: tensor})
    rboxes, scores, classes = postprocess(outs, names, opt.score_thres,
                                          opt.nms_thres, opt.agnostic_nms)
    print(f"Forward+post: {(time()-t0)*1000:.1f} ms")

    if len(rboxes) == 0:
        print("No detections.")
        cv2.imwrite(opt.output, img0)
        return

    rboxes = rboxes[:opt.max_det]
    scores = scores[:opt.max_det]
    classes = classes[:opt.max_det]

    rboxes = scale_back(rboxes, r, pad, img0.shape[:2])
    print(f"Found {len(rboxes)} oriented objects.")
    for (cx, cy, w, h, th), sc, cid in zip(rboxes, scores, classes):
        name = DOTA_CLASSES[cid] if cid < len(DOTA_CLASSES) else f"cls{cid}"
        print(f"  {name:20s} conf={sc:.2f} cx={cx:.1f} cy={cy:.1f} "
              f"w={w:.1f} h={h:.1f} theta={math.degrees(th):+.1f}")

    draw(img0, rboxes, scores, classes)
    cv2.imwrite(opt.output, img0)
    print(f"Saved: {opt.output}")


if __name__ == "__main__":
    main()