#!/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).
    """
    orig_h, orig_w = image.shape[:2]
    m_h, m_w = input_size

    scale = min(m_h / orig_h, m_w / orig_w)
    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)

    # Left-top aligned padding
    input_bgr = cv2.copyMakeBorder(
        img_resized, 0, m_h - new_h, 0, m_w - new_w,
        cv2.BORDER_CONSTANT, value=(114, 114, 114)
    )

    input_rgb = cv2.cvtColor(input_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    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 to 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. If reg_max given and channels match, apply DFL first."""
    if reg_max is not None and bbox_preds.shape[-1] == 4 * reg_max:
        bbox_preds = dfl_decode(bbox_preds, reg_max)
    lt = bbox_preds[:, :2]
    rb = bbox_preds[:, 2:]
    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 (left-top aligned)."""
    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 process_mask(protos, mask_coeffs, boxes, model_shape, orig_shape, scale):
    """
    Process masks from prototypes and coefficients.

    Args:
        protos: (nm, proto_h, proto_w) mask prototypes
        mask_coeffs: (N, nm) mask coefficients
        boxes: (N, 4) boxes in model coordinates (xyxy)
        model_shape: (h, w) model input shape
        orig_shape: (h, w) original image shape
        scale: preprocessing scale factor

    Returns:
        masks: (N, orig_h, orig_w) binary masks
    """
    nm, proto_h, proto_w = protos.shape
    N = len(mask_coeffs)

    if N == 0:
        return np.zeros((0, orig_shape[0], orig_shape[1]), dtype=np.uint8)

    # Compute masks: (N, nm) @ (nm, proto_h*proto_w) -> (N, proto_h, proto_w)
    masks = (mask_coeffs @ protos.reshape(nm, -1)).reshape(N, proto_h, proto_w)
    # Sigmoid activation
    masks = 1 / (1 + np.exp(-masks))

    # Scale boxes to proto coordinates
    width_ratio = proto_w / model_shape[1]
    height_ratio = proto_h / model_shape[0]
    boxes_proto = boxes.copy()
    boxes_proto[:, [0, 2]] *= width_ratio
    boxes_proto[:, [1, 3]] *= height_ratio

    # Crop masks to bounding box regions
    for i, (x1, y1, x2, y2) in enumerate(boxes_proto.astype(int)):
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(proto_w, x2), min(proto_h, y2)
        masks[i, :y1, :] = 0
        masks[i, y2:, :] = 0
        masks[i, :, :x1] = 0
        masks[i, :, x2:] = 0

    # Upsample masks to model input size
    masks_upsampled = np.zeros((N, model_shape[0], model_shape[1]), dtype=np.float32)
    for i in range(N):
        masks_upsampled[i] = cv2.resize(masks[i], (model_shape[1], model_shape[0]),
                                        interpolation=cv2.INTER_LINEAR)

    # Crop to valid region (left-top aligned) and resize to original
    new_h = int(orig_shape[0] * scale)
    new_w = int(orig_shape[1] * scale)
    masks_cropped = masks_upsampled[:, :new_h, :new_w]

    # Resize to original image size
    masks_final = np.zeros((N, orig_shape[0], orig_shape[1]), dtype=np.uint8)
    for i in range(N):
        mask_resized = cv2.resize(masks_cropped[i], (orig_shape[1], orig_shape[0]),
                                  interpolation=cv2.INTER_LINEAR)
        masks_final[i] = (mask_resized > 0.5).astype(np.uint8)

    return masks_final


def main():
    parser = argparse.ArgumentParser(description='YOLOv8-Seg ONNX Inference')
    parser.add_argument('-m', '--model', type=str, default='yolov8n-seg_640x640.onnx',
                        dest='model_path', help='Path to YOLOv8-Seg *.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_seg.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})

    # Parse outputs: 9 detection outputs + 1 proto
    strides = [8, 16, 32]
    conf_raw = -np.log(1 / opt.score_thres - 1)
    detections = []
    all_mask_coeffs = []

    # Process each scale
    for scale_idx, stride in enumerate(strides):
        box_idx = scale_idx * 3
        cls_idx = scale_idx * 3 + 1
        mask_idx = scale_idx * 3 + 2

        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)
        mask_data = outputs[mask_idx]  # (1, H, W, nm)

        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
        box_data = box_data[0].reshape(-1, box_channels)
        cls_data = cls_data[0].reshape(-1, cls_data.shape[-1])
        mask_data = mask_data[0].reshape(-1, mask_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
        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]
        v_mask = mask_data[valid_mask]

        # Generate anchors
        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]])
            all_mask_coeffs.append(v_mask[i])

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

    detections = np.array(detections)
    all_mask_coeffs = np.array(all_mask_coeffs)

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

    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]
    all_mask_coeffs = all_mask_coeffs[indices]

    # Get proto output (last output)
    proto = outputs[-1][0]  # (nm, proto_h, proto_w)

    # Process masks (before scaling boxes!)
    boxes_model = detections[:, :4].copy()
    masks = process_mask(proto, all_mask_coeffs, boxes_model, imgsz, orig_shape, scale)

    # Scale boxes 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 i, det in enumerate(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)

        # Draw mask
        if i < len(masks):
            mask = masks[i]
            mask_colored = np.zeros_like(img0)
            mask_colored[mask > 0] = color
            img0 = cv2.addWeighted(img0, 1.0, mask_colored, 0.5, 0)

        # Draw box
        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()