"""
eval_bie_yolo11m.py — Kneron BIE / ONNX / NEF Evaluator for YOLOv11m
======================================================================
Usage examples
--------------
# Evaluate BIE  (sigmoid was removed before BIE creation → always pass --apply-sigmoid)
python eval_bie_yolo11m.py \
    --bie     ./models/model.bie \
    --img-dir ./dataset/val/images \
    --lbl-dir ./dataset/val/labels \
    --iou-thres 0.45 \
    --apply-sigmoid

# Evaluate no-sigmoid ONNX  (optimized_best_no_sigmoid.onnx)
python eval_bie_yolo11m.py \
    --onnx-path ./models/optimized_best_no_sigmoid.onnx \
    --img-dir   ./dataset/val/images \
    --lbl-dir   ./dataset/val/labels \
    --iou-thres 0.45 \
    --apply-sigmoid

# Evaluate original ONNX that still has sigmoid baked in
python eval_bie_yolo11m.py \
    --onnx-path ./models/best.onnx \
    --img-dir   ./dataset/val/images \
    --lbl-dir   ./dataset/val/labels \
    --iou-thres 0.45
    # DO NOT pass --apply-sigmoid here — sigmoid is inside the model already

# Sweep NMS IoU to find the best threshold
python eval_bie_yolo11m.py \
    --bie     ./models/model.bie \
    --img-dir ./dataset/val/images \
    --lbl-dir ./dataset/val/labels \
    --iou-sweep \
    --apply-sigmoid
"""

import argparse
import os
import time
import numpy as np
import cv2

from pathlib import Path
from collections import defaultdict

try:
    import ktc
except ImportError:
    print("ERROR: ktc (Kneron Toolchain) is not installed in this environment.")
    exit(1)


# ==========================================
# ARGUMENT PARSING
# ==========================================
def parse_args():
    parser = argparse.ArgumentParser(
        description="YOLOv11m Kneron BIE / ONNX / NEF Evaluator",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Model source — exactly one required
    model_group = parser.add_mutually_exclusive_group(required=True)
    model_group.add_argument("--bie",       metavar="PATH", help="Path to .bie file")
    model_group.add_argument("--onnx-path", metavar="PATH", help="Path to .onnx file")
    model_group.add_argument("--nef",       metavar="PATH", help="Path to .nef file")

    # Dataset
    parser.add_argument("--img-dir",    required=True, metavar="PATH",
                        help="Directory containing validation images (.jpg / .png)")
    parser.add_argument("--lbl-dir",    required=True, metavar="PATH",
                        help="Directory containing YOLO-format .txt label files")

    # Kneron platform / model metadata
    parser.add_argument("--platform",   default="730", metavar="ID",
                        help="Kneron platform ID (e.g. 730, 520)")
    parser.add_argument("--input-name", default="images", metavar="NAME",
                        help="ONNX input tensor name")
    parser.add_argument("--imgsz",      type=int, default=640,
                        help="Input image size (square)")

    # Detection thresholds
    parser.add_argument("--conf-thres", type=float, default=0.25,
                        help="Minimum confidence to keep a detection")
    parser.add_argument("--score-margin", type=float, default=0.0,
                        help="Minimum gap between top-1 and top-2 class scores. "
                             "0.0 = disabled (recommended starting point). "
                             "Only increase if you see many duplicate/noisy class predictions.")
    parser.add_argument("--max-det",    type=int, default=300,
                        help="Maximum detections returned per image")

    # ── IMPORTANT FLAG ───────────────────────────────────────────────────────
    # Pass --apply-sigmoid whenever the model does NOT contain a sigmoid node:
    #   • BIE / NEF built from a no-sigmoid ONNX  → ALWAYS pass this flag
    #   • Optimized ONNX with sigmoid removed      → pass this flag
    #   • Original ONNX with sigmoid still inside  → do NOT pass this flag
    # ─────────────────────────────────────────────────────────────────────────
    parser.add_argument("--apply-sigmoid", action="store_true", default=False,
                        help="Apply sigmoid to raw class logits in post-processing. "
                             "Required for BIE/NEF where sigmoid was removed before compilation.")

    # Visualisation
    parser.add_argument("--vis-dir", default=None, metavar="PATH",
                        help="If set, saves annotated images for the first 20 samples")

    # NMS IoU — one of these is required
    iou_group = parser.add_mutually_exclusive_group(required=True)
    iou_group.add_argument("--iou-thres", type=float, metavar="FLOAT",
                           help="Fixed NMS IoU threshold (e.g. 0.45)")
    iou_group.add_argument("--iou-sweep", action="store_true", default=False,
                           help="Sweep NMS IoU from 0.005 to 0.995 in 0.005 steps "
                                "and report the best mAP@0.50")

    args = parser.parse_args()

    # Warn clearly when evaluating BIE/NEF without --apply-sigmoid — almost always a mistake
    if (args.bie or args.nef) and not args.apply_sigmoid:
        print("\n" + "!" * 60)
        print("  WARNING: You are evaluating a BIE/NEF without --apply-sigmoid.")
        print("  If the sigmoid was removed before BIE creation (the standard")
        print("  Kneron workflow), raw logits will be used as confidence scores")
        print("  and mAP will drop by 20-30%.")
        print("  Add --apply-sigmoid unless you are certain your BIE still")
        print("  contains a sigmoid node.")
        print("!" * 60 + "\n")

    return args


# ==========================================
# PREPROCESSING
# ==========================================
def letterbox(img, new_size=640, color=(114, 114, 114)):
    """
    Resize-with-padding (letterbox) using OpenCV.
    Returns: (padded_img, scale, pad_left, pad_top)
    """
    h, w   = img.shape[:2]
    scale  = new_size / max(h, w)
    new_w  = int(round(w * scale))
    new_h  = int(round(h * scale))

    img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    pad_w  = new_size - new_w
    pad_h  = new_size - new_h
    top    = pad_h // 2;    bottom = pad_h - top
    left   = pad_w // 2;    right  = pad_w - left

    img_padded = cv2.copyMakeBorder(
        img_resized, top, bottom, left, right,
        cv2.BORDER_CONSTANT, value=color,
    )
    return img_padded, scale, left, top


def preprocess_image(img_path, img_size):
    """
    Load an image, apply letterbox, convert BGR→RGB, normalise to [0, 1].
    Returns:
        input_tensor : float32 ndarray  (1, 3, img_size, img_size)  range [0, 1]
        orig_img     : uint8  ndarray   (H0, W0, 3)  original BGR image for visualisation
        w0, h0       : original image dimensions
        scale        : resize scale factor applied
        pad_left     : horizontal padding (pixels)
        pad_top      : vertical   padding (pixels)
    """
    img = cv2.imread(str(img_path))
    if img is None:
        raise ValueError(f"Could not read image: {img_path}")

    h0, w0 = img.shape[:2]
    img_lb, scale, pad_left, pad_top = letterbox(img, img_size)

    # BGR → RGB, HWC → CHW, uint8 → float32 [0, 1]
    img_input = img_lb[:, :, ::-1].transpose(2, 0, 1)
    img_input = np.expand_dims(img_input, 0).astype(np.float32) / 255.0

    return img_input, img, w0, h0, scale, pad_left, pad_top


# ==========================================
# NMS
# ==========================================
def nms_boxes(boxes, scores, iou_thres):
    """Pure-NumPy per-class NMS. boxes: (N,4) xyxy, scores: (N,)"""
    if len(boxes) == 0:
        return np.array([], dtype=np.int64)

    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas  = (x2 - x1) * (y2 - y1)
    order  = scores.argsort()[::-1]
    keep   = []

    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        ovr   = inter / (areas[i] + areas[order[1:]] - inter + 1e-7)
        order = order[np.where(ovr <= iou_thres)[0] + 1]

    return np.array(keep, dtype=np.int64)


# ==========================================
# POST-PROCESSING  (YOLOv11m)
# ==========================================
def postprocess_yolo11(
    raw_tensor,
    scale, pad_left, pad_top,
    w0, h0,
    conf_thres,
    iou_thres,
    apply_sigmoid,
    score_margin,
    max_det,
    max_nms=30000,
):
    """
    YOLOv11m post-processing pipeline
    ----------------------------------
    1. Normalise tensor shape  →  (4+nc, N)
    2. Sigmoid on class scores (only when model has no sigmoid node)
    3. Confidence pre-filter   (conf > conf_thres)
    4. Score-margin filter     (optional — disabled by default)
    5. cx,cy,w,h  →  x1,y1,x2,y2
    6. Un-letterbox + clip to original image space
    7. Degenerate-box filter   (< 2 px on either side)
    8. Per-class NMS
    9. Cap to max_det

    Returns
    -------
    list of (class_id: int, confidence: float, box: [x1,y1,x2,y2])
    sorted by confidence descending.
    """
    # ── 1. Normalise shape ────────────────────────────────────────────────────
    t = raw_tensor
    if   t.ndim == 4: t = t[0, 0]   # (1,1,C,N) → (C,N)
    elif t.ndim == 3: t = t[0]       # (1,C,N)   → (C,N)
    # else already (C,N) or (N,C)

    if t.shape[0] > t.shape[1]:
        t = t.T                       # ensure (C, N)  where C = 4+nc

    boxes_raw        = t[:4]          # (4, N)
    class_scores_raw = t[4:]          # (nc, N)

    # ── 2. Sigmoid (when sigmoid was stripped from the model) ─────────────────
    if apply_sigmoid:
        class_scores = 1.0 / (1.0 + np.exp(-np.clip(class_scores_raw, -88.0, 88.0)))
    else:
        class_scores = class_scores_raw

    # ── 3. Confidence pre-filter ──────────────────────────────────────────────
    max_cls_scores  = class_scores.max(axis=0)          # (N,)
    candidate_mask  = max_cls_scores > conf_thres
    if not candidate_mask.any():
        return []

    boxes_cxywh  = boxes_raw[:, candidate_mask].T       # (K, 4)
    class_scores = class_scores[:, candidate_mask].T    # (K, nc)

    # ── 4. Score-margin filter (default: disabled via score_margin=0.0) ───────
    if score_margin > 0.0 and class_scores.shape[1] > 1:
        sorted_s     = np.sort(class_scores, axis=1)
        margin_mask  = (sorted_s[:, -1] - sorted_s[:, -2]) > score_margin
        boxes_cxywh  = boxes_cxywh[margin_mask]
        class_scores = class_scores[margin_mask]
        if len(boxes_cxywh) == 0:
            return []

    # ── 5. cx,cy,w,h → x1,y1,x2,y2 ──────────────────────────────────────────
    cx, cy, bw, bh = (boxes_cxywh[:, i] for i in range(4))
    x1 = cx - bw / 2;  x2 = cx + bw / 2
    y1 = cy - bh / 2;  y2 = cy + bh / 2
    boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)    # (K, 4)

    # ── 6. Un-letterbox → original image space ────────────────────────────────
    boxes_orig       = np.empty_like(boxes_xyxy)
    boxes_orig[:, 0] = np.clip((boxes_xyxy[:, 0] - pad_left) / scale, 0, w0)
    boxes_orig[:, 1] = np.clip((boxes_xyxy[:, 1] - pad_top)  / scale, 0, h0)
    boxes_orig[:, 2] = np.clip((boxes_xyxy[:, 2] - pad_left) / scale, 0, w0)
    boxes_orig[:, 3] = np.clip((boxes_xyxy[:, 3] - pad_top)  / scale, 0, h0)

    class_ids   = np.argmax(class_scores, axis=1).astype(np.int32)
    confidences = class_scores[np.arange(len(class_scores)), class_ids]

    # ── 7. Degenerate-box filter ──────────────────────────────────────────────
    valid       = ((boxes_orig[:, 2] - boxes_orig[:, 0]) >= 2) & \
                  ((boxes_orig[:, 3] - boxes_orig[:, 1]) >= 2)
    boxes_orig  = boxes_orig[valid]
    confidences = confidences[valid]
    class_ids   = class_ids[valid]

    if len(boxes_orig) == 0:
        return []

    # Cap before NMS to avoid O(N²) explosion on very noisy outputs
    if len(boxes_orig) > max_nms:
        top_idx     = np.argsort(-confidences)[:max_nms]
        boxes_orig  = boxes_orig[top_idx]
        confidences = confidences[top_idx]
        class_ids   = class_ids[top_idx]

    # ── 8. Per-class NMS ──────────────────────────────────────────────────────
    final_dets = []
    for cls in np.unique(class_ids):
        mask = (class_ids == cls)
        keep = nms_boxes(boxes_orig[mask], confidences[mask], iou_thres)
        for i in keep:
            final_dets.append((
                int(cls),
                float(confidences[mask][i]),
                boxes_orig[mask][i].tolist(),
            ))

    # ── 9. Sort by confidence, cap to max_det ─────────────────────────────────
    final_dets.sort(key=lambda x: x[1], reverse=True)
    return final_dets[:max_det]


# ==========================================
# GROUND-TRUTH LOADING
# ==========================================
def load_ground_truth(txt_path, img_w, img_h):
    """Load YOLO-format labels and convert to absolute xyxy pixel coordinates."""
    gts = []
    if not os.path.exists(txt_path):
        return gts
    with open(txt_path, "r") as f:
        for line in f:
            parts  = line.strip().split()
            if len(parts) < 5:
                continue
            cls_id      = int(parts[0])
            cx, cy, w, h = map(float, parts[1:5])
            x1 = (cx - w / 2) * img_w;  y1 = (cy - h / 2) * img_h
            x2 = (cx + w / 2) * img_w;  y2 = (cy + h / 2) * img_h
            gts.append((cls_id, [x1, y1, x2, y2]))
    return gts


# ==========================================
# EVALUATION METRICS
# ==========================================
def compute_iou(boxA, boxB):
    xA = max(boxA[0], boxB[0]);  yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2]);  yB = min(boxA[3], boxB[3])
    inter = max(0.0, xB - xA) * max(0.0, yB - yA)
    areaA = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
    areaB = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
    union = areaA + areaB - inter
    return inter / union if union > 0 else 0.0


def match_detections_to_gt(dets, gts, iou_t):
    """
    Greedy match detections → ground-truth boxes (sorted by confidence desc).
    Returns list of (tp_flag: int, confidence: float, class_id: int).
    """
    matched_gts = set()
    results     = []
    for cls_p, conf_p, box_p in sorted(dets, key=lambda x: -x[1]):
        best_iou = 0.0;  best_j = -1
        for j, (cls_g, box_g) in enumerate(gts):
            if cls_p != cls_g or j in matched_gts:
                continue
            iou = compute_iou(box_p, box_g)
            if iou > best_iou:
                best_iou = iou;  best_j = j
        if best_iou >= iou_t and best_j >= 0:
            results.append((1, conf_p, cls_p))
            matched_gts.add(best_j)
        else:
            results.append((0, conf_p, cls_p))
    return results


def compute_ap_per_class(tp_list, conf_list, n_gt):
    """Compute AP using 101-point interpolation (COCO style)."""
    if n_gt == 0:
        return 0.0, np.zeros(101), np.zeros(101)
    idx       = np.argsort(-np.array(conf_list))
    tp_sorted = np.array(tp_list)[idx]
    tp_cum    = np.cumsum(tp_sorted)
    fp_cum    = np.cumsum(1 - tp_sorted)
    recall    = tp_cum    / (n_gt + 1e-16)
    precision = tp_cum    / (tp_cum + fp_cum + 1e-16)
    recall    = np.concatenate(([0.0], recall,    [1.0]))
    precision = np.concatenate(([1.0], precision, [0.0]))
    precision = np.flip(np.maximum.accumulate(np.flip(precision)))
    r_interp  = np.linspace(0, 1, 101)
    p_interp  = np.interp(r_interp, recall, precision)
    return float(np.mean(p_interp)), p_interp, r_interp


# ==========================================
# VISUALISATION (optional)
# ==========================================
def draw_predictions(img, dets, gts, save_path):
    """Green = predictions, Red = ground truth."""
    for cls_id, conf, (x1, y1, x2, y2) in dets:
        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
        cv2.putText(img, f"cls{cls_id} {conf:.2f}",
                    (int(x1), max(int(y1) - 5, 12)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 0), 1)
    for cls_id, (x1, y1, x2, y2) in gts:
        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
    cv2.imwrite(save_path, img)


# ==========================================
# INFERENCE WRAPPER
# ==========================================
def run_inference(args, input_tensor):
    """Call ktc.kneron_inference with the correct source file."""
    if args.onnx_path:
        return ktc.kneron_inference(
            [input_tensor],
            onnx_file=args.onnx_path,
            input_names=[args.input_name],
        )
    elif args.bie:
        return ktc.kneron_inference(
            [input_tensor],
            bie_file=args.bie,
            input_names=[args.input_name],
            platform=int(args.platform),
        )
    else:  # nef
        return ktc.kneron_inference(
            [input_tensor],
            nef_file=args.nef,
            input_names=[args.input_name],
            platform=int(args.platform),
        )


# ==========================================
# MAIN
# ==========================================
def main():
    args       = parse_args()
    model_name = Path(args.bie or args.onnx_path or args.nef).name

    # ── NMS IoU sweep values ──────────────────────────────────────────────────
    if args.iou_sweep:
        nms_iou_sweep = np.round(np.arange(0.005, 1.000, 0.005), 3)
    else:
        nms_iou_sweep = np.array([round(args.iou_thres, 3)])

    # ── Header ────────────────────────────────────────────────────────────────
    print("\n" + "=" * 60)
    print(f"  MODEL         : {model_name}")
    print(f"  ARCHITECTURE  : YOLOv11m")
    print("=" * 60)
    print(f"  imgsz         : {args.imgsz}")
    print(f"  conf_thres    : {args.conf_thres}")
    print(f"  score_margin  : {args.score_margin}  {'(disabled)' if args.score_margin == 0.0 else ''}")
    print(f"  apply_sigmoid : {args.apply_sigmoid}")
    print(f"  max_det       : {args.max_det}")
    if args.iou_sweep:
        print(f"  NMS IoU sweep : {nms_iou_sweep[0]:.3f} → {nms_iou_sweep[-1]:.3f}  "
              f"({len(nms_iou_sweep)} steps)")
    else:
        print(f"  NMS IoU thres : {args.iou_thres}")
    print("=" * 60 + "\n")

    # ── Dataset ───────────────────────────────────────────────────────────────
    images = sorted([
        f for f in os.listdir(args.img_dir)
        if f.lower().endswith((".jpg", ".jpeg", ".png"))
    ])
    if not images:
        print("ERROR: No images found in --img-dir.")
        return

    if args.vis_dir:
        os.makedirs(args.vis_dir, exist_ok=True)

    # ── Accumulators ─────────────────────────────────────────────────────────
    coco_iou_thresholds = np.round(np.arange(0.50, 0.96, 0.05), 2)

    sweep_results = {
        iou: {
            "class_detections": {t: defaultdict(list) for t in coco_iou_thresholds},
            "det_counts":       [],
        }
        for iou in nms_iou_sweep
    }
    class_gt_count = defaultdict(int)
    latencies      = []

    # ── Inference loop ────────────────────────────────────────────────────────
    print(f"Running inference on {len(images)} images...\n")

    for idx, img_name in enumerate(images):
        img_path = os.path.join(args.img_dir, img_name)
        lbl_path = os.path.join(args.lbl_dir, img_name.rsplit(".", 1)[0] + ".txt")

        input_tensor, orig_img, w0, h0, scale, pad_left, pad_top = \
            preprocess_image(img_path, args.imgsz)

        # Inference
        t0      = time.perf_counter()
        outputs = run_inference(args, input_tensor)
        latencies.append((time.perf_counter() - t0) * 1000.0)

        raw_output = outputs[0]

        # Ground truth
        gts = load_ground_truth(lbl_path, w0, h0)
        for cls_g, _ in gts:
            class_gt_count[cls_g] += 1

        # Post-process for every NMS IoU candidate
        for nms_iou in nms_iou_sweep:
            dets = postprocess_yolo11(
                raw_output,
                scale, pad_left, pad_top, w0, h0,
                args.conf_thres,
                nms_iou,
                args.apply_sigmoid,
                args.score_margin,
                args.max_det,
            )
            sweep_results[nms_iou]["det_counts"].append(len(dets))

            # Visualise first 20 images
            if args.vis_dir and idx < 20:
                vis_subdir = os.path.join(args.vis_dir, f"iou_{nms_iou:.3f}")
                os.makedirs(vis_subdir, exist_ok=True)
                draw_predictions(
                    orig_img.copy(), dets, gts,
                    os.path.join(vis_subdir, img_name),
                )

            # Accumulate per-class TP/FP at every COCO IoU threshold
            for coco_iou in coco_iou_thresholds:
                for tp_flag, conf, cls_p in match_detections_to_gt(dets, gts, coco_iou):
                    sweep_results[nms_iou]["class_detections"][coco_iou][cls_p].append(
                        (tp_flag, conf)
                    )

        # Per-image progress
        if args.iou_sweep:
            all_det_counts = [sweep_results[iou]["det_counts"][-1] for iou in nms_iou_sweep]
            print(f"  [{idx+1:4d}/{len(images)}] {img_name:<35s} "
                  f"gts={len(gts):3d}  dets=[{min(all_det_counts)}-{max(all_det_counts)}]")
        else:
            n_dets = sweep_results[nms_iou_sweep[0]]["det_counts"][-1]
            print(f"  [{idx+1:4d}/{len(images)}] {img_name:<35s} "
                  f"gts={len(gts):3d}  dets={n_dets:3d}")

    # ==========================================
    # RESULTS
    # ==========================================
    all_classes = sorted(class_gt_count.keys())

    def compute_ap_matrix(nms_iou):
        ap_mat = np.zeros((len(coco_iou_thresholds), len(all_classes)))
        for ci, coco_iou in enumerate(coco_iou_thresholds):
            for ki, cls in enumerate(all_classes):
                dets_cls = sweep_results[nms_iou]["class_detections"][coco_iou].get(cls, [])
                if not dets_cls:
                    continue
                ap, _, _ = compute_ap_per_class(
                    [d[0] for d in dets_cls],
                    [d[1] for d in dets_cls],
                    class_gt_count[cls],
                )
                ap_mat[ci, ki] = ap
        return ap_mat

    avg_latency = float(np.mean(latencies))

    if args.iou_sweep:
        print("\n" + "=" * 80)
        print("  NMS IoU SWEEP RESULTS")
        print("=" * 80)
        best_map50  = -1.0
        best_iou_50 = None
        for nms_iou in nms_iou_sweep:
            ap_matrix = compute_ap_matrix(nms_iou)
            map50     = float(np.mean(ap_matrix[0, :]))
            avg_dets  = float(np.mean(sweep_results[nms_iou]["det_counts"]))
            if map50 > best_map50:
                best_map50  = map50
                best_iou_50 = nms_iou
            print(f"  NMS IoU: {nms_iou:>6.3f}  |  mAP@0.50: {map50:>7.4f}  |  "
                  f"avg dets/img: {avg_dets:>6.1f}")
        print()
        print(f"  ★ Best mAP@0.50 : {best_map50:.4f}  at NMS IoU = {best_iou_50:.3f}")
        print(f"  Avg Latency     : {avg_latency:.2f} ms")
        print("=" * 80)

    else:
        ap_matrix   = compute_ap_matrix(nms_iou_sweep[0])
        map50       = float(np.mean(ap_matrix[0, :]))
        map50_95    = float(np.mean(ap_matrix))

        print("\n" + "=" * 60)
        print("  RESULTS")
        print("=" * 60)
        print(f"  mAP @ 0.50      : {map50:.4f}   ({map50 * 100:.2f}%)")
        print(f"  mAP @ 0.50:0.95 : {map50_95:.4f}   ({map50_95 * 100:.2f}%)")
        print(f"  Avg Latency     : {avg_latency:.2f} ms")
        print()

        # Per-class breakdown
        print("  Per-class AP @ 0.50:")
        for ki, cls in enumerate(all_classes):
            ap_cls = float(ap_matrix[0, ki])
            n_gt   = class_gt_count[cls]
            print(f"    class {cls:>3d} : {ap_cls:.4f}  ({n_gt} gt boxes)")
        print("=" * 60)


if __name__ == "__main__":
    main()