import os
import cv2
import math
import numpy as np
from time import time
import argparse
import logging
import axengine as ort

logging.basicConfig(
    level=logging.DEBUG,
    format='[%(name)s] [%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s',
    datefmt='%H:%M:%S'
)
logger = logging.getLogger("YOLO26-OBB")

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):
    """Model input (NCHW or NHWC) -> H, W, layout name."""
    shape = list(shape)
    if len(shape) == 4 and shape[-1] == 3:
        h = int(shape[1] or 1024)
        w = int(shape[2] or 1024)
        return h, w, "NHWC"
    if len(shape) == 4 and shape[1] == 3:
        h = int(shape[2] or 1024)
        w = int(shape[3] or 1024)
        return h, w, "NCHW"
    return 1024, 1024, "NCHW"


def preprocess_image(image, input_size=(1024, 1024), layout="NCHW", padding_value=114):
    """LetterBox + BGR→RGB; uint8 NHWC or NCHW (same geometry as onnx_infer letterbox)."""
    orig_h, orig_w = image.shape[:2]
    new_h, new_w = input_size
    r = min(new_h / orig_h, new_w / orig_w)
    new_unpad_w = round(orig_w * r)
    new_unpad_h = round(orig_h * r)
    dw = new_w - new_unpad_w
    dh = new_h - new_unpad_h
    dw /= 2.0
    dh /= 2.0
    if (orig_w, orig_h) != (new_unpad_w, new_unpad_h):
        image = cv2.resize(image, (new_unpad_w, new_unpad_h), interpolation=cv2.INTER_LINEAR)
    top = round(dh - 0.1)
    bottom = round(dh + 0.1)
    left = round(dw - 0.1)
    right = round(dw + 0.1)
    padded = cv2.copyMakeBorder(
        image, top, bottom, left, right,
        cv2.BORDER_CONSTANT, value=(padding_value,) * 3,
    )
    ratio_pad = (r, (left, top))
    rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
    if layout == "NHWC":
        tensor = rgb[None, ...].astype(np.uint8)
    else:
        tensor = np.transpose(rgb, (2, 0, 1))[None, ...].astype(np.uint8)
    return tensor, ratio_pad, (orig_h, orig_w)


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


def dfl_decode(box_pred, reg_max):
    N = box_pred.shape[0]
    box_pred = box_pred.reshape(N, 4, reg_max)
    box_pred = softmax(box_pred, axis=-1)
    proj = np.arange(reg_max, dtype=np.float32)
    return np.sum(box_pred * proj, axis=-1)


def decode_obb(box_preds, angle_preds, anchors, stride, reg_max=None):
    """dist2rbox -> xywhr in pixels."""
    if reg_max is not None and box_preds.shape[-1] == 4 * reg_max and reg_max > 1:
        box_preds = dfl_decode(box_preds, reg_max)

    angle = angle_preds.reshape(-1)
    cos_a = np.cos(angle)
    sin_a = np.sin(angle)

    lt = box_preds[:, :2]
    rb = box_preds[:, 2:]
    xf = (rb[:, 0] - lt[:, 0]) * 0.5
    yf = (rb[:, 1] - lt[:, 1]) * 0.5

    cx = xf * cos_a - yf * sin_a + anchors[:, 0]
    cy = xf * sin_a + yf * cos_a + anchors[:, 1]
    w = lt[:, 0] + rb[:, 0]
    h = lt[:, 1] + rb[:, 1]
    return np.stack([cx * stride, cy * stride, w * stride, h * stride, angle], axis=1)


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


def batch_probiou(obb1, obb2, eps=1e-7):
    x1 = obb1[:, 0:1]
    y1 = obb1[:, 1:2]
    x2 = obb2[:, 0][None, :]
    y2 = obb2[:, 1][None, :]
    a1, b1, c1 = (v[:, None] for v in _get_covariance_matrix(obb1))
    a2_full, b2_full, c2_full = _get_covariance_matrix(obb2)
    a2 = a2_full[None, :]
    b2 = b2_full[None, :]
    c2 = c2_full[None, :]

    sum_ab = (a1 + a2) * (b1 + b2) - (c1 + c2) ** 2
    t1 = ((a1 + a2) * (y1 - y2) ** 2 + (b1 + b2) * (x1 - x2) ** 2) / (sum_ab + eps) * 0.25
    t2 = ((c1 + c2) * (x2 - x1) * (y1 - y2)) / (sum_ab + eps) * 0.5
    inner = ((a1 * b1 - c1 ** 2).clip(min=0.0) * (a2 * b2 - c2 ** 2).clip(min=0.0))
    t3 = np.log(sum_ab / (4.0 * np.sqrt(inner) + 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 nms_rotated_probiou(rboxes, scores, classes, iou_thres, max_wh=7680.0, agnostic=False):
    """Fast NMS with probiou."""
    if rboxes.size == 0:
        return np.empty((0,), dtype=np.int64)
    boxes = rboxes.copy()
    if not agnostic:
        offset = classes.astype(np.float32) * float(max_wh)
        boxes[:, 0] = boxes[:, 0] + offset
        boxes[:, 1] = boxes[:, 1] + offset
    order = np.argsort(-scores)
    sorted_boxes = boxes[order]
    ious = batch_probiou(sorted_boxes, sorted_boxes)
    n = sorted_boxes.shape[0]
    triu = np.triu(np.ones((n, n), dtype=bool), k=1)
    ious = ious * triu
    keep_mask = (ious >= iou_thres).sum(axis=0) <= 0
    return order[keep_mask]


def scale_rboxes_lefttop(rboxes, ratio_pad, orig_shape):
    """Letterbox inverse for xywhr centers."""
    rboxes = rboxes.copy()
    gain, (pad_x, pad_y) = ratio_pad
    rboxes[:, 0] -= pad_x
    rboxes[:, 1] -= pad_y
    rboxes[:, :4] /= gain
    rboxes[:, 0] = np.clip(rboxes[:, 0], 0, orig_shape[1])
    rboxes[:, 1] = np.clip(rboxes[:, 1], 0, orig_shape[0])
    return rboxes


def regularize_rbox(rboxes):
    """Theta in [0, pi/2); swap w/h when needed."""
    rboxes = rboxes.copy()
    t_mod = np.mod(rboxes[:, 4], np.pi)
    swap = t_mod >= (np.pi / 2.0)
    if np.any(swap):
        w_old = rboxes[swap, 2].copy()
        rboxes[swap, 2] = rboxes[swap, 3]
        rboxes[swap, 3] = w_old
    rboxes[:, 4] = np.mod(rboxes[:, 4], np.pi / 2.0)
    return rboxes


def rbox_to_corners(rbox):
    cx, cy, w, h, ag = rbox
    cos_a, sin_a = math.cos(ag), math.sin(ag)
    wx, wy = w / 2.0 * cos_a, w / 2.0 * sin_a
    hx, hy = -h / 2.0 * sin_a, h / 2.0 * cos_a
    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 main():
    ap = argparse.ArgumentParser(description='YOLO26-OBB Inference (AXERARuntime)')
    ap.add_argument('--model-path', type=str, default='yolo26n-obb_1024x1024.axmodel')
    ap.add_argument('--test-img', type=str, default='boats.jpg')
    ap.add_argument('--img-save-path', type=str, default='result_yolo26_obb.jpg')
    ap.add_argument('--score-thres', type=float, default=0.25)
    ap.add_argument('--nms-thres', type=float, default=0.45)
    ap.add_argument('--num-classes', type=int, default=15)
    ap.add_argument('--max-det', type=int, default=300)
    ap.add_argument('--agnostic-nms', action='store_true', help='Class-agnostic NMS.')
    ap.add_argument('--providers', type=str, default='AxEngineExecutionProvider')
    opt = ap.parse_args()

    if not os.path.exists(opt.model_path):
        logger.error(f"Model not found: {opt.model_path}")
        return
    if not os.path.exists(opt.test_img):
        logger.error(f"Image not found: {opt.test_img}")
        return

    t0 = time()
    providers = [p.strip() for p in opt.providers.split(",") if p.strip()] or None
    sess = ort.InferenceSession(opt.model_path, providers=providers)
    logger.debug(f"\033[1;31mLoad model time = {(time() - t0) * 1000:.2f} ms\033[0m")

    inp = sess.get_inputs()[0]
    input_name = inp.name
    m_h, m_w, layout = infer_hw_layout(inp.shape)

    img0 = cv2.imread(opt.test_img)
    if img0 is None:
        logger.error(f"Cannot read image: {opt.test_img}")
        return

    t0 = time()
    input_tensor, ratio_pad, orig_shape = preprocess_image(
        img0.copy(), input_size=(m_h, m_w), layout=layout
    )
    logger.debug(f"\033[1;31mPre-process time = {(time() - t0) * 1000:.2f} ms\033[0m")

    t0 = time()
    output_names = [o.name for o in sess.get_outputs()]
    outputs = sess.run(output_names, {input_name: input_tensor})
    logger.debug(f"\033[1;31mForward time = {(time() - t0) * 1000:.2f} ms\033[0m")

    t0 = time()
    strides = [8, 16, 32]
    conf_raw = -math.log(1.0 / opt.score_thres - 1.0)
    rboxes_all = []
    scores_all = []
    classes_all = []

    for scale_idx, stride in enumerate(strides):
        box_idx = scale_idx * 3
        cls_idx = scale_idx * 3 + 1
        ang_idx = scale_idx * 3 + 2
        if ang_idx >= len(outputs):
            continue

        box_data = outputs[box_idx]
        cls_data = outputs[cls_idx]
        ang_data = outputs[ang_idx]

        H, W = box_data.shape[1:3]
        box_channels = box_data.shape[-1]
        reg_max = None
        if box_channels > 4 and box_channels % 4 == 0:
            reg_max = box_channels // 4

        box_data = box_data[0].reshape(-1, box_channels)
        cls_data = cls_data[0].reshape(-1, cls_data.shape[-1])
        ang_data = ang_data[0].reshape(-1, ang_data.shape[-1])

        if cls_data.shape[-1] == 1:
            cls_logits = cls_data[:, 0]
            cls_ids = np.zeros(len(cls_logits), dtype=np.int32)
        else:
            cls_logits = np.max(cls_data, axis=1)
            cls_ids = np.argmax(cls_data, axis=1)

        valid = cls_logits >= conf_raw
        if not np.any(valid):
            continue

        v_box = box_data[valid]
        v_ang = ang_data[valid]
        v_score = 1.0 / (1.0 + np.exp(-cls_logits[valid]))
        v_id = cls_ids[valid]

        gy, gx = np.indices((H, W))
        anchors = np.stack((gx.ravel(), gy.ravel()), axis=-1).astype(np.float32) + 0.5
        anchors = anchors[valid]

        rboxes = decode_obb(v_box, v_ang, anchors, stride, reg_max)
        rboxes_all.append(rboxes)
        scores_all.append(v_score)
        classes_all.append(v_id)

    if len(rboxes_all) == 0:
        logger.info("No detections found.")
        cv2.imwrite(opt.img_save_path, img0)
        return

    rboxes_all = np.concatenate(rboxes_all, axis=0).astype(np.float32)
    scores_all = np.concatenate(scores_all, axis=0).astype(np.float32)
    classes_all = np.concatenate(classes_all, axis=0).astype(np.int32)

    keep = nms_rotated_probiou(
        rboxes_all, scores_all, classes_all,
        iou_thres=opt.nms_thres, agnostic=opt.agnostic_nms,
    )
    keep = keep[: opt.max_det]
    logger.debug(f"\033[1;31mPost-process time = {(time() - t0) * 1000:.2f} ms\033[0m")

    if len(keep) == 0:
        logger.info("No detections after NMS.")
        cv2.imwrite(opt.img_save_path, img0)
        return

    final_rboxes = rboxes_all[keep]
    final_scores = scores_all[keep]
    final_classes = classes_all[keep]

    final_rboxes = scale_rboxes_lefttop(final_rboxes, ratio_pad, orig_shape)
    final_rboxes = regularize_rbox(final_rboxes)

    logger.info(f"\033[1;32mDraw Results ({len(final_rboxes)} oriented objects):\033[0m")
    for i in range(len(final_rboxes)):
        cx, cy, w, h, theta = final_rboxes[i]
        conf = float(final_scores[i])
        cid = int(final_classes[i])
        name = DOTA_CLASSES[cid] if cid < len(DOTA_CLASSES) else f"cls{cid}"
        color = DOTA_COLORS[cid % len(DOTA_COLORS)]
        logger.info(
            f"  {name:20s} conf={conf:.2f} cx={cx:.1f} cy={cy:.1f} "
            f"w={w:.1f} h={h:.1f} theta={math.degrees(theta):+.1f}deg"
        )
        corners = rbox_to_corners((cx, cy, w, h, theta)).astype(np.int32)
        cv2.polylines(img0, [corners], isClosed=True, color=color, thickness=2, lineType=cv2.LINE_AA)
        label = f"{name} {conf:.2f}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        x_text, y_text = int(corners[0][0]), max(0, int(corners[0][1]) - 5)
        cv2.rectangle(img0, (x_text, y_text - th - 2), (x_text + tw + 2, y_text + 2), color, -1)
        cv2.putText(img0, label, (x_text + 1, y_text - 1),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)

    cv2.imwrite(opt.img_save_path, img0)
    logger.info(f"Saved to {opt.img_save_path}")


if __name__ == "__main__":
    main()