"""Minimal runner for YOLO11 inference + mAP using helper utilities.

This script is intentionally small: preprocessing, postprocessing and
evaluation live in `helper.py`. Use the static paths at the top.
"""

import glob
import os
import time
from collections import defaultdict

import cv2
import kp
import numpy as np
from helper import (
    preprocess_image,
    postprocess_yolo11,
    evaluate_map,
    load_gt_labels,
    _merge_nodes,
)


# ─── CONFIG ───────────────────────────────────────────────

NEF_PATH      = "/home/milan/milanop/kneron-model-conversion-od/models11m_wuthoutrobo/models_730.nef"
FW_PATH       = "/home/milan/Downloads/kneron_plus_v3.0.0/kneron_plus/res/firmware/KL730/kp_firmware.tar"
IMAGE_FOLDER  = "/home/milan/milanop/kneron-model-conversion-od/20_crop/train/images"
OUTPUT_FOLDER = "/home/milan/milanop/kneron-model-conversion-od/20_crop/train/output_yolo11/"
LABELS_FOLDER = "/home/milan/milanop/kneron-model-conversion-od/20_crop/train/labels/"

CLASSES = [
 "helmet", "head", "gloves", "shoes",
    "hands", "foot", "safety-vest", "glasses", "face", "facemask",
]

# CLASSES = [
#     "face","face-mask" ,"foot", "glasses","gloves","helmet","hands","head","shoes"
# ,"safety-suit","safety-vest"]
NUM_CLASSES = len(CLASSES)   # 10

IOU_MAP       = 0.50   # mAP evaluation IoU threshold
NMS_IOU       = 0.040   # NMS IoU threshold (set higher for better mAP, lower for better precision)
APPLY_SIGMOID = True   # set False if NEF already bakes sigmoid into output
CONF_THRESHOLD = 0.25
SCORE_MARGIN = NMS_IOU
MAX_DET = 300

def draw_detections(img, detections, gt_boxes):
    vis = img.copy()
    for (cls_id, x1, y1, x2, y2) in gt_boxes:
        cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
    for (cls_id, score, bbox) in detections:
        x1, y1, x2, y2 = [int(v) for v in bbox]
        label = f"{CLASSES[cls_id]} {score:.2f}"
        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv2.putText(vis, label, (x1, max(y1 - 5, 0)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
    return vis


def save_detections(label_path, detections):
    with open(label_path, "w", encoding="utf-8") as f:
        for cls_id, score, bbox in detections:
            x1, y1, x2, y2 = bbox
            f.write(
                f"{cls_id} {CLASSES[cls_id]} {score:.6f} "
                f"{x1:.2f} {y1:.2f} {x2:.2f} {y2:.2f}\n"
            )

def run_inference(device, model_id, image_path,
                  conf_thres, nms_iou, score_margin, max_det):
    """
    Full pipeline for one image:
        preprocess_image (CPU letterbox)
        → KP SDK send/receive  (resize+padding DISABLED — we already did it)
        → retrieve float nodes
        → _merge_nodes
        → postprocess_yolo11

    IMPORTANT: KP_RESIZE_DISABLE + KP_PADDING_DISABLE are set because
    preprocess_image already produced a 640×640 letterboxed image.
    Leaving resize/padding enabled would apply a second transform and
    make pad_left / pad_top completely wrong.
    """
    # ── Preprocess (CPU-side letterbox) ──────────────────────────────────
    try:
        # helper.preprocess_image now returns img_lb as well so we can
        # reuse it for RGB565 conversion without re-running letterbox.
        img_input, img, img_lb, w0, h0, scale, pad_left, pad_top = preprocess_image(
            image_path, img_size=640
        )
    except ValueError as e:
        print(f"  WARNING: {e}")
        return None, [], 0.0

    # ── USB buffer guard (RGB565 = 2 bytes/px, limit ~4 MB) ──────────────
    # After letterbox the image is always 640×640 = 819 200 px = 1.6 MB RGB565,
    # so this guard only fires if someone changes img_size above 1448.
    MAX_RGB565_BYTES = 4 * 1024 * 1024
    if 640 * 640 * 2 > MAX_RGB565_BYTES:
        print("  SKIP (letterboxed image exceeds USB buffer)")
        return None, [], 0.0

    # ── Ensure even dimensions on original image (KP_ERROR_IMAGE_ODD_WIDTH_23)
    # The letterboxed canvas is always 640×640 (even), so this check covers
    # any edge case where img_size was set to an odd number.
    lb_h, lb_w = 640, 640   # letterbox output is always square

    # ── Convert letterboxed BGR image (returned by preprocess) to RGB565
    img_rgb565 = cv2.cvtColor(img_lb, cv2.COLOR_BGR2BGR565)

    # ── Build KP inference descriptor ────────────────────────────────────
    infer_input = kp.GenericImageInferenceDescriptor(
        model_id=model_id,
        inference_number=0,
        input_node_image_list=[
            kp.GenericInputNodeImage(
                image=img_rgb565,
                image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
                # ↓↓ CRITICAL: we already letterboxed CPU-side ↓↓
                resize_mode=kp.ResizeMode.KP_RESIZE_DISABLE,
                padding_mode=kp.PaddingMode.KP_PADDING_DISABLE,
                normalize_mode=kp.NormalizeMode.KP_NORMALIZE_YOLO,
            )
        ],
    )

    # ── Send / receive ────────────────────────────────────────────────────
    t0 = time.time()
    kp.inference.generic_image_inference_send(
        device_group=device,
        generic_inference_input_descriptor=infer_input,
    )
    result = kp.inference.generic_image_inference_receive(device)
    t_ms = (time.time() - t0) * 1000.0

    # ── Collect all output nodes ──────────────────────────────────────────
    # Single merged node  → shape (1, 4+NC, 8400)  for 640-input
    # Three FPN nodes     → shapes (1, 4+NC, 6400), (…1600), (…400)
    # SDK raises exception when node_idx is out of range — use as stop signal.
    raw_nodes = []
    node_idx  = 0
    while True:
        try:
            node_out = kp.inference.generic_inference_retrieve_float_node(
                node_idx=node_idx,
                generic_raw_result=result,
                channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW,
            )
            arr = node_out.ndarray
            raw_nodes.append(arr)
            print(f"  node[{node_idx}] shape={arr.shape}  dtype={arr.dtype}"
                  f"  min={arr.min():.3f}  max={arr.max():.3f}")
            node_idx += 1
        except Exception:
            break

    if not raw_nodes:
        print("  ERROR: no output nodes retrieved!")
        return img, [], t_ms

    # ── Merge multi-head nodes → single (C, N_total) tensor ──────────────
    merged = _merge_nodes(raw_nodes)
    print(f"  merged (C,N)={merged.shape}  expected C={4 + NUM_CLASSES}")

    dets = postprocess_yolo11(
        raw_tensor   = merged,
        scale        = scale,
        pad_left     = pad_left,
        pad_top      = pad_top,
        w0           = w0,
        h0           = h0,
        conf_thres   = conf_thres,
        iou_thres    = nms_iou,
        apply_sigmoid= APPLY_SIGMOID,
        score_margin = score_margin,
        max_det      = max_det,
    )
    return img, dets, t_ms


def main():
    os.makedirs(OUTPUT_FOLDER, exist_ok=True)
    labels_output_dir = os.path.join(OUTPUT_FOLDER, "labels")
    os.makedirs(labels_output_dir, exist_ok=True)

    print("Connecting to KL730 …")
    device = kp.core.connect_devices(usb_port_ids=[0])
    print("Connected!")

    print("Loading firmware …")
    kp.core.load_firmware_from_file(
        device_group=device,
        scpu_fw_path=FW_PATH,
        ncpu_fw_path=FW_PATH,
    )
    print("Firmware loaded!")

    print("Loading NEF model …")
    model_desc = kp.core.load_model_from_file(
        device_group=device,
        file_path=NEF_PATH,
    )
    model_id = model_desc.models[0].id
    print(f"Model loaded!  ID: {model_id}\n")

    image_paths = []
    for ext in ("*.jpg", "*.jpeg", "*.png", "*.bmp"):
        image_paths += glob.glob(os.path.join(IMAGE_FOLDER, ext))
    image_paths.sort()
    print(f"Found {len(image_paths)} images in {IMAGE_FOLDER}\n")

    all_preds   = {}
    all_gts     = {}
    infer_times = []
    t_total     = time.time()

    for idx, image_path in enumerate(image_paths):
        fname      = os.path.basename(image_path)
        stem       = os.path.splitext(fname)[0]
        label_path = os.path.join(LABELS_FOLDER, stem + ".txt")
        out_path   = os.path.join(OUTPUT_FOLDER, fname)
        pred_path  = os.path.join(labels_output_dir, stem + ".txt")

        # GT uses original image dims — read them before preprocess_image
        # (preprocess_image also returns w0/h0, but we need gt_list before
        #  run_inference to keep the loop clean)
        img_check = cv2.imread(str(image_path))
        if img_check is None:
            continue
        h0_gt, w0_gt = img_check.shape[:2]
        gt_list = load_gt_labels(label_path, w0_gt, h0_gt)

        img, dets, t_ms = run_inference(
            device, model_id, image_path,
            conf_thres=CONF_THRESHOLD,
            nms_iou=NMS_IOU,
            score_margin=SCORE_MARGIN,
            max_det=MAX_DET,
        )
        if img is None:
            continue

        infer_times.append(t_ms)
        all_preds[idx] = [(c, s, *b) for (c, s, b) in dets]
        all_gts[idx]   = gt_list

        vis = draw_detections(
            img, dets,
            [(c, x1, y1, x2, y2) for (c, x1, y1, x2, y2) in gt_list]
        )
        cv2.imwrite(out_path, vis)
        save_detections(pred_path, dets)

        print(f"[{idx+1}/{len(image_paths)}] {fname}")
        print(f"  Detections : {len(dets)}")
        print(f"  GT boxes   : {len(gt_list)}")
        print(f"  Infer time : {t_ms:.1f} ms")
        print(f"  Saved      : {out_path}")
        print(f"  Labels     : {pred_path}")
        print("─" * 60)

    # ── Timing summary ────────────────────────────────────────────────────
    if infer_times:
        total = time.time() - t_total
        print(f"\n{'='*60}")
        print(f"  Total images   : {len(infer_times)}")
        print(f"  Total time     : {total:.2f} s")
        print(f"  Avg infer      : {np.mean(infer_times):.1f} ms")
        print(f"  Min infer      : {np.min(infer_times):.1f} ms")
        print(f"  Max infer      : {np.max(infer_times):.1f} ms")
        print(f"  Throughput     : {len(infer_times)/total:.2f} FPS")
        print(f"{'='*60}")

    # ── mAP evaluation ────────────────────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"  mAP evaluation  (IoU threshold = {IOU_MAP:.2f})")
    print(f"{'='*60}")

    per_class_ap, map_score = evaluate_map(all_preds, all_gts, NUM_CLASSES, IOU_MAP)

    col_w = max(len(c) for c in CLASSES) + 2
    print(f"  {'Class':<{col_w}}  {'AP@50':>8}  {'#GT':>6}")
    print(f"  {'-'*col_w}  {'-------':>8}  {'----':>6}")
    for cls_id, cls_name in enumerate(CLASSES):
        ap   = per_class_ap[cls_id]
        n_gt = sum(1 for gts in all_gts.values() for (c, *_) in gts if c == cls_id)
        ap_s = f"{ap*100:6.2f}%" if not np.isnan(ap) else "    N/A"
        print(f"  {cls_name:<{col_w}}  {ap_s:>8}  {n_gt:>6}")
    print(f"  {'-'*col_w}  {'-------':>8}")
    print(f"  {'mAP@50':<{col_w}}  {map_score*100:6.2f}%")
    print(f"{'='*60}\n")

    kp.core.disconnect_devices(device_group=device)
    print("Device disconnected cleanly.")


if __name__ == "__main__":
    main()