#!/usr/bin/env python3
import onnxruntime as ort
import cv2
import numpy as np
import argparse
import os


def preprocess_image(image, input_size=(640, 640)):
    """
    Preprocess image with left-top aligned letterbox (same as official YOLO).

    Args:
        image: BGR image (H, W, C)
        input_size: (height, width) target size

    Returns:
        input_tensor: (1, 3, H, W) float32 tensor normalized to [0, 1]
        scale: scale ratio used for resizing
        original_shape: (orig_h, orig_w) original image shape
    """
    orig_h, orig_w = image.shape[:2]
    m_h, m_w = input_size

    # Calculate scale (keep aspect ratio)
    scale = min(m_h / orig_h, m_w / orig_w)

    # Resize
    new_w, new_h = int(orig_w * scale), int(orig_h * scale)
    img_resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    # Pad to target size (left-top aligned)
    input_bgr = cv2.copyMakeBorder(
        img_resized, 0, m_h - new_h, 0, m_w - new_w,
        cv2.BORDER_CONSTANT, value=(114, 114, 114)
    )

    # BGR -> RGB, normalize to [0, 1]
    input_rgb = cv2.cvtColor(input_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0

    # HWC -> CHW, add batch dimension
    input_tensor = np.transpose(input_rgb, (2, 0, 1))[None, ...]

    return input_tensor, scale, (orig_h, orig_w)


def softmax(x, axis=-1):
    """Compute softmax along axis."""
    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=16):
    """
    Decode DFL (Distribution Focal Loss) box predictions.

    Args:
        box_pred: (N, 4 * reg_max) raw box predictions
        reg_max: number of DFL bins

    Returns:
        (N, 4) decoded ltrb distances
    """
    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)  # (N, 4)


def decode_bboxes(bbox_preds, anchors, stride, reg_max=None):
    """
    Decode bounding boxes from predictions.

    Args:
        bbox_preds: (N, 4) or (N, 4*reg_max) - distance predictions in ltrb format
        anchors: (N, 2) - anchor points (x, y) with offset 0.5
        stride: scalar - stride value
        reg_max: if not None, apply DFL decoding first

    Returns:
        boxes: (N, 4) in xyxy format
    """
    if reg_max is not None and bbox_preds.shape[-1] == 4 * reg_max:
        bbox_preds = dfl_decode(bbox_preds, reg_max)

    # dist2bbox: ltrb to xyxy
    lt = bbox_preds[:, :2]  # left, top
    rb = bbox_preds[:, 2:]  # right, bottom

    x1y1 = anchors - lt
    x2y2 = anchors + rb

    boxes = np.hstack([x1y1, x2y2]) * stride
    return boxes


def scale_boxes_lefttop(boxes, scale, orig_shape):
    """Scale boxes from model output to original image coordinates."""
    boxes = boxes.copy()
    boxes[..., :4] /= scale
    boxes[..., [0, 2]] = np.clip(boxes[..., [0, 2]], 0, orig_shape[1])
    boxes[..., [1, 3]] = np.clip(boxes[..., [1, 3]], 0, orig_shape[0])
    return boxes


def main():
    parser = argparse.ArgumentParser(description='YOLOv8-Det ONNX Inference')
    parser.add_argument('-m', '--model', type=str, default='yolov8n_640x640.onnx',
                        dest='model_path', help='Path to YOLOv8 Detection *.onnx Model.')
    parser.add_argument('-i', '--img', type=str, default='bus.jpg',
                        dest='test_img', help='Path to Test Image.')
    parser.add_argument('-o', '--output', type=str, default='result_yolov8_det.jpg',
                        dest='img_save_path', help='Path to Save Result Image.')
    parser.add_argument('--score-thres', type=float, default=0.25,
                        help='Confidence threshold.')
    parser.add_argument('--nms-thres', type=float, default=0.7,
                        help='IoU threshold for NMS.')
    opt = parser.parse_args()

    if not os.path.exists(opt.model_path):
        print(f"Error: Model not found: {opt.model_path}")
        return

    if not os.path.exists(opt.test_img):
        print(f"Error: Image not found: {opt.test_img}")
        return

    # Load ONNX model
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    try:
        session = ort.InferenceSession(opt.model_path, providers=providers)
    except:
        session = ort.InferenceSession(opt.model_path, providers=['CPUExecutionProvider'])

    input_name = session.get_inputs()[0].name
    output_names = [o.name for o in session.get_outputs()]
    input_shape = session.get_inputs()[0].shape
    imgsz = (input_shape[2], input_shape[3])

    # Load image
    img0 = cv2.imread(opt.test_img)
    if img0 is None:
        print(f"Error: Cannot read image: {opt.test_img}")
        return

    # Preprocess
    img, scale, orig_shape = preprocess_image(img0.copy(), imgsz)

    # Inference
    outputs = session.run(output_names, {input_name: img})

    # Post-process
    strides = [8, 16, 32]
    conf_raw = -np.log(1 / opt.score_thres - 1)
    detections = []

    # Process each scale (6 outputs: 2 outputs per scale [box, cls])
    for scale_idx, stride in enumerate(strides):
        box_idx = scale_idx * 2
        cls_idx = scale_idx * 2 + 1

        if box_idx >= len(outputs) or cls_idx >= len(outputs):
            continue

        box_data = outputs[box_idx]  # (1, H, W, C) where C = 4 or 4*reg_max
        cls_data = outputs[cls_idx]  # (1, H, W, nc)

        H, W = box_data.shape[1:3]
        box_channels = box_data.shape[-1]

        # Determine if DFL is used (YOLOv8: 4*reg_max = 64)
        reg_max = None
        if box_channels > 4 and box_channels % 4 == 0:
            reg_max = box_channels // 4

        # Reshape to (H*W, ...)
        box_data = box_data[0].reshape(-1, box_channels)
        cls_data = cls_data[0].reshape(-1, cls_data.shape[-1])

        # Get max class scores
        cls_scores = np.max(cls_data, axis=1)
        cls_ids = np.argmax(cls_data, axis=1)

        # Filter by confidence (on raw logits)
        valid_mask = cls_scores >= conf_raw
        if not np.any(valid_mask):
            continue

        v_box = box_data[valid_mask]
        v_score = 1 / (1 + np.exp(-cls_scores[valid_mask]))
        v_id = cls_ids[valid_mask]

        # Generate anchors for this scale (grid indices + 0.5)
        gy, gx = np.indices((H, W))
        anchors = np.stack((gx.ravel(), gy.ravel()), axis=-1).astype(np.float32) + 0.5
        anchors = anchors[valid_mask]

        # Decode boxes
        boxes = decode_bboxes(v_box, anchors, stride, reg_max)

        # Store detections
        for i in range(len(boxes)):
            detections.append([*boxes[i], v_score[i], v_id[i]])

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

    detections = np.array(detections)

    # NMS
    xywh = detections[:, :4].copy()
    xywh[:, 2] = xywh[:, 2] - xywh[:, 0]  # w
    xywh[:, 3] = xywh[:, 3] - xywh[:, 1]  # h

    indices = cv2.dnn.NMSBoxes(xywh.tolist(), detections[:, 4].tolist(), opt.score_thres, opt.nms_thres)

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

    indices = indices.flatten()
    detections = detections[indices]

    # Scale to original image
    detections[:, :4] = scale_boxes_lefttop(detections[:, :4], scale, orig_shape)

    # COCO class names
    coco_names = [
        "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
        "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
        "scissors", "teddy bear", "hair drier", "toothbrush"
    ]

    # Official Ultralytics colors (BGR format)
    base_colors = [
        (255, 42, 4), (235, 219, 11), (243, 243, 243), (183, 223, 0), (104, 31, 17),
        (221, 111, 255), (79, 68, 255), (0, 237, 204), (68, 243, 0), (255, 0, 189),
        (255, 180, 0), (186, 0, 221), (255, 255, 0), (0, 192, 38), (179, 255, 1),
        (255, 36, 125), (104, 0, 123), (108, 27, 255), (47, 109, 252), (11, 255, 162),
    ]

    for det in detections:
        box = det[:4].astype(int)
        conf = det[4]
        cls_id = int(det[5])
        color = [int(c) for c in base_colors[cls_id % len(base_colors)]]
        cls_name = coco_names[cls_id] if cls_id < len(coco_names) else str(cls_id)

        cv2.rectangle(img0, (box[0], box[1]), (box[2], box[3]), color, 2)
        label = f"{cls_name} {conf:.2f}"
        cv2.putText(img0, label, (box[0], box[1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    cv2.imwrite(opt.img_save_path, img0)
    print(f"Done! Found {len(detections)} objects. Result saved to {opt.img_save_path}")


if __name__ == "__main__":
    main()