"""Helper module: single preprocessing, single postprocessing, and mAP.

Expose:
 - preprocess_image(img_path, img_size=640) -> img_input, img, img_lb, w0, h0, scale, pad_left, pad_top
 - postprocess_yolo11(raw_tensor, scale, pad_left, pad_top, w0, h0, conf_thres, iou_thres, apply_sigmoid, score_margin, max_det)
 - evaluate_map(all_preds, all_gts, num_classes, iou_thresh=0.50)
 - load_gt_labels(label_path, img_w, img_h)
"""
import os
from collections import defaultdict

import cv2
import numpy as np


def _letterbox(img, new_size=640, color=(114, 114, 114)):
    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=640):
    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], add batch dim
    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, img_lb, w0, h0, scale, pad_left, pad_top


def _nms_boxes(boxes, scores, iou_thres):
    if len(boxes) == 0:
        return []
    x1 = boxes[:, 0]; y1 = boxes[:, 1]
    x2 = boxes[:, 2]; y2 = 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)
        inds = np.where(ovr <= iou_thres)[0]
        order = order[inds + 1]
    return np.array(keep)


def _merge_nodes(raw_nodes):
    parts = []
    for raw in raw_nodes:
        r = raw
        while r.ndim > 3 and r.shape[0] == 1:
            r = r[0]
        if r.ndim == 3:
            if r.shape[0] == 1:
                r = r[0]
            else:
                r = r.reshape(r.shape[0], -1)
        if r.ndim == 2 and r.shape[0] > r.shape[1]:
            r = r.T
        parts.append(r)
    merged = np.concatenate(parts, axis=1)
    return merged


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):
    raw = (raw_tensor[0, 0] if raw_tensor.ndim == 4
           else raw_tensor[0] if raw_tensor.ndim == 3
           else raw_tensor)
    if raw.shape[0] > raw.shape[1]:
        raw = raw.T
    class_scores_raw = raw[4:]
    class_scores = (1.0 / (1.0 + np.exp(-np.clip(class_scores_raw, -88.0, 88.0)))
                    if apply_sigmoid else class_scores_raw)
    max_cls_scores = class_scores.max(axis=0)
    candidate_mask = max_cls_scores > conf_thres
    if not candidate_mask.any():
        return []
    boxes_cxywh = raw[:4, candidate_mask].T
    class_scores = class_scores[:, candidate_mask].T
    if score_margin > 0.0 and class_scores.shape[1] > 1:
        sorted_scores = np.sort(class_scores, axis=1)
        margin_mask = (sorted_scores[:, -1] - sorted_scores[:, -2]) > score_margin
        boxes_cxywh = boxes_cxywh[margin_mask]
        class_scores = class_scores[margin_mask]
    if len(boxes_cxywh) == 0:
        return []
    class_ids = np.argmax(class_scores, axis=1).astype(np.int32)
    confidences = class_scores[np.arange(len(class_scores)), class_ids]
    x1 = boxes_cxywh[:, 0] - boxes_cxywh[:, 2] / 2
    y1 = boxes_cxywh[:, 1] - boxes_cxywh[:, 3] / 2
    x2 = boxes_cxywh[:, 0] + boxes_cxywh[:, 2] / 2
    y2 = boxes_cxywh[:, 1] + boxes_cxywh[:, 3] / 2
    boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)
    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)
    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) > 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]
    final_dets = []
    for cls in np.unique(class_ids):
        cls_mask = (class_ids == cls)
        keep = _nms_boxes(boxes_orig[cls_mask], confidences[cls_mask], iou_thres)
        for idx in keep:
            final_dets.append((int(cls), float(confidences[cls_mask][idx]), boxes_orig[cls_mask][idx].tolist()))
    return sorted(final_dets, key=lambda x: x[1], reverse=True)[:max_det]


def box_iou(box1, box2):
    a1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
    a2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
    ix1 = np.maximum(box1[:, None, 0], box2[None, :, 0])
    iy1 = np.maximum(box1[:, None, 1], box2[None, :, 1])
    ix2 = np.minimum(box1[:, None, 2], box2[None, :, 2])
    iy2 = np.minimum(box1[:, None, 3], box2[None, :, 3])
    inter = np.maximum(ix2 - ix1, 0) * np.maximum(iy2 - iy1, 0)
    return inter / (a1[:, None] + a2[None, :] - inter + 1e-9)


def load_gt_labels(label_path, img_w, img_h):
    boxes = []
    if not os.path.exists(label_path):
        return boxes
    with open(label_path) 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
            boxes.append((cls_id, x1, y1, x2, y2))
    return boxes


def compute_ap(recalls, precisions):
    ap = 0.0
    for t in np.linspace(0, 1, 11):
        p = precisions[recalls >= t]
        ap += np.max(p) if p.size > 0 else 0.0
    return ap / 11.0


def evaluate_map(all_preds, all_gts, num_classes, iou_thresh=0.50):
    per_class_ap = []
    for cls in range(num_classes):
        detections = [
            (score, img_id, det_idx)
            for img_id, preds in all_preds.items()
            for det_idx, (c, score, *_) in enumerate(preds)
            if c == cls
        ]
        detections.sort(key=lambda x: -x[0])
        n_gt = sum(1 for gts in all_gts.values() for (c, *_) in gts if c == cls)
        if n_gt == 0:
            per_class_ap.append(float("nan"))
            continue
        if not detections:
            per_class_ap.append(0.0)
            continue
        matched = defaultdict(set)
        tp = np.zeros(len(detections))
        fp = np.zeros(len(detections))
        for rank, (score, img_id, det_idx) in enumerate(detections):
            pred = all_preds[img_id][det_idx]
            pred_box = np.array([[pred[2], pred[3], pred[4], pred[5]]])
            gt_cands = [
                (gi, gt)
                for gi, gt in enumerate(all_gts.get(img_id, []))
                if gt[0] == cls and gi not in matched[img_id]
            ]
            if not gt_cands:
                fp[rank] = 1
                continue
            gt_arr = np.array([[g[1], g[2], g[3], g[4]] for _, g in gt_cands])
            ious = box_iou(pred_box, gt_arr)[0]
            best = int(np.argmax(ious))
            if ious[best] >= iou_thresh:
                tp[rank] = 1
                matched[img_id].add(gt_cands[best][0])
            else:
                fp[rank] = 1
        cum_tp = np.cumsum(tp)
        cum_fp = np.cumsum(fp)
        recalls = cum_tp / (n_gt + 1e-9)
        precisions = cum_tp / (cum_tp + cum_fp + 1e-9)
        per_class_ap.append(compute_ap(recalls, precisions))
    valid = [ap for ap in per_class_ap if not np.isnan(ap)]
    return per_class_ap, float(np.mean(valid)) if valid else 0.0
