"""YOLO 共享后处理 —— 按 (algorithm_type, task_type) 分发

支持版本：yolo5 / yolo8 / yolo11 / yolo26
支持任务：detect / segment / classify / pose / obb

输出结构差异：
- yolo5 detect:  output [1, N, 5+nc]  行格式 [cx, cy, w, h, obj, cls_scores...]
- yolo8/11/26 detect: output [1, nc+4, N]  列格式 [cx, cy, w, h, cls_score_0, cls_score_1, ...]  无 obj
- yolo5 segment: outputs[0] [1, N, 5+nc+nm], outputs[1] [1, nm, mh, mw] (protos)
- yolo8/11/26 segment: outputs[0] [1, nc+4+nm, N], outputs[1] [1, nm, mh, mw]
- classify (all): output [1, nc]  top-1 logits
- yolo8/11/26 pose: output [1, nc+4+nk*3, N]  (nk 通常 17)
- yolo8/11/26 obb:  output [1, nc+4+1, N]  末列是角度 (radians)

本模块仅依赖 numpy + opencv，被 OnnxEngine / OpenVinoEngine 复用。
YoloPytorchEngine 直接用 ultralytics 的 Results 对象，不走本模块。
"""

import logging
import numpy as np

logger = logging.getLogger("analysis.engines.yolo_post")

try:
    import cv2
    _CV2 = True
except Exception:
    cv2 = None
    _CV2 = False


def decode_outputs(outputs, algorithm_type, task_type, labels, input_size,
                   conf_threshold, iou_threshold, orig_size, num_classes=None):
    """主入口：返回 list[dict]，每条至少含 box/label/score，按任务附加上下文。

    outputs: list[np.ndarray]  原始模型输出（顺序与 onnx session.get_outputs() 一致）
    """
    try:
        if task_type == "classify":
            return _decode_classify(outputs, labels)
        if task_type == "detect":
            return _decode_detect(outputs, algorithm_type, labels, input_size,
                                  conf_threshold, iou_threshold, orig_size, num_classes)
        if task_type == "segment":
            return _decode_segment(outputs, algorithm_type, labels, input_size,
                                   conf_threshold, iou_threshold, orig_size, num_classes)
        if task_type == "pose":
            return _decode_pose(outputs, algorithm_type, labels, input_size,
                                conf_threshold, iou_threshold, orig_size, num_classes)
        if task_type == "obb":
            return _decode_obb(outputs, algorithm_type, labels, input_size,
                               conf_threshold, iou_threshold, orig_size, num_classes)
        logger.warning("yolo_post: 未知任务类型 %s", task_type)
        return []
    except Exception as e:
        logger.warning("yolo_post decode err: %s", e)
        return []


def _is_v5(algorithm_type):
    return algorithm_type in ("yolo5", "yolov5", "v5")


def _xywh2xyxy(xywh):
    x, y, w, h = xywh[..., 0], xywh[..., 1], xywh[..., 2], xywh[..., 3]
    return np.stack([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)


def _nms(boxes, scores, iou_threshold):
    if len(boxes) == 0:
        return []
    boxes_list = boxes.tolist() if hasattr(boxes, "tolist") else [list(b) for b in boxes]
    scores_list = scores.tolist() if hasattr(scores, "tolist") else list(scores)
    if _CV2:
        try:
            idx = cv2.dnn.NMSBoxes(boxes_list, scores_list, 0.0, float(iou_threshold))
            if idx is None or len(idx) == 0:
                return []
            if hasattr(idx, "flatten"):
                idx = idx.flatten()
            return [int(i) for i in idx]
        except Exception as e:
            logger.debug("cv2 NMS fallback: %s", e)
    # OpenCV 不可用或 NMS 失败时的朴素实现
    order = np.argsort(-np.asarray(scores_list, dtype=np.float32))
    suppressed = np.zeros(len(scores_list), dtype=bool)
    keep = []
    boxes_arr = np.asarray(boxes_list, dtype=np.float32)
    for i in order:
        if suppressed[i]:
            continue
        keep.append(int(i))
        for j in order:
            if j == i or suppressed[j]:
                continue
            if _iou(boxes_arr[i], boxes_arr[j]) > float(iou_threshold):
                suppressed[j] = True
    return keep


def _iou(a, b):
    xa, ya = max(a[0], b[0]), max(a[1], b[1])
    xb, yb = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0.0, xb - xa), max(0.0, yb - ya)
    inter = iw * ih
    area_a = max(0.0, a[2] - a[0]) * max(0.0, a[3] - a[1])
    area_b = max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1])
    union = area_a + area_b - inter
    return inter / union if union > 0 else 0.0


def _scale_box(box, sx, sy, orig_w, orig_h):
    x1 = max(0, min(orig_w - 1, box[0] * sx))
    y1 = max(0, min(orig_h - 1, box[1] * sy))
    x2 = max(0, min(orig_w - 1, box[2] * sx))
    y2 = max(0, min(orig_h - 1, box[3] * sy))
    return [int(x1), int(y1), int(x2), int(y2)]


# ===================== classify =====================
def _decode_classify(outputs, labels):
    out = np.asarray(outputs[0])
    if out.ndim == 2 and out.shape[0] == 1:
        out = out[0]
    if out.ndim != 1:
        # 某些导出会带 batch 维 + 多余维度
        out = out.reshape(-1)
    cid = int(np.argmax(out))
    score = float(out[cid])
    label = labels[cid] if 0 <= cid < len(labels) else str(cid)
    return [{"box": [0, 0, 0, 0], "label": str(label), "score": score, "task": "classify"}]


# ===================== detect =====================
def _decode_detect(outputs, algorithm_type, labels, input_size, conf, iou, orig_size, nc=None):
    arr = np.asarray(outputs[0])
    iw, ih = input_size
    orig_w, orig_h = orig_size
    sx, sy = float(orig_w) / iw, float(orig_h) / ih
    if _is_v5(algorithm_type):
        # [1, N, 5+nc] 或 [N, 5+nc]
        if arr.ndim == 3 and arr.shape[0] == 1:
            arr = arr[0]
        if arr.ndim != 2 or arr.shape[1] < 6:
            return []
        nc = arr.shape[1] - 5 if nc is None else nc
        if nc < 1:
            nc = 1
        objs = arr[:, 4]
        cls_scores = arr[:, 5:5 + nc]
        # score = obj * cls
        scores_all = objs[:, None] * cls_scores  # [N, nc]
        best_cls = np.argmax(scores_all, axis=1)
        best_score = scores_all[np.arange(len(arr)), best_cls]
        mask = best_score >= conf
        if not np.any(mask):
            return []
        xywh = arr[mask, :4]
        boxes_raw = _xywh2xyxy(xywh)
        scores = best_score[mask]
        clses = best_cls[mask]
    else:
        # yolo8/11/26: [1, nc+4, N]  需转置
        if arr.ndim == 3 and arr.shape[0] == 1:
            arr = arr[0]
        if arr.ndim != 2:
            return []
        # 行格式可能是 [4+nc, N] —— 转成 [N, 4+nc]
        if arr.shape[0] < arr.shape[1] and arr.shape[0] >= 4:
            arr = arr.T
        if arr.ndim != 2 or arr.shape[1] < 5:
            return []
        nc = arr.shape[1] - 4 if nc is None else nc
        if nc < 1:
            nc = 1
        boxes_raw = _xywh2xyxy(arr[:, :4])
        cls_scores = arr[:, 4:4 + nc]
        best_cls = np.argmax(cls_scores, axis=1)
        best_score = cls_scores[np.arange(len(arr)), best_cls]
        mask = best_score >= conf
        if not np.any(mask):
            return []
        boxes_raw = boxes_raw[mask]
        scores = best_score[mask]
        clses = best_cls[mask]

    boxes = [_scale_box(b, sx, sy, orig_w, orig_h) for b in boxes_raw]
    keep = _nms(boxes, scores, iou)
    out = []
    for i in keep:
        label = labels[int(clses[i])] if 0 <= int(clses[i]) < len(labels) else str(int(clses[i]))
        out.append({"box": boxes[i], "label": str(label), "score": float(scores[i]), "task": "detect"})
    return out


# ===================== segment =====================
def _decode_segment(outputs, algorithm_type, labels, input_size, conf, iou, orig_size, nc=None):
    arr = np.asarray(outputs[0])
    proto = np.asarray(outputs[1]) if len(outputs) > 1 else None
    iw, ih = input_size
    orig_w, orig_h = orig_size
    sx, sy = float(orig_w) / iw, float(orig_h) / ih
    if _is_v5(algorithm_type):
        if arr.ndim == 3 and arr.shape[0] == 1:
            arr = arr[0]
        if arr.ndim != 2 or arr.shape[1] < 6:
            return []
        nm = arr.shape[1] - 5 - (nc or (arr.shape[1] - 5))
        # 推断 nm：若 nc 已知
        if nc is None:
            # 退化：用 arr.shape[1]-5 / 2 估计（不可靠），优先按 nc 推断
            nm = 32
            nc = arr.shape[1] - 5 - nm
        else:
            nm = arr.shape[1] - 5 - nc
        if nm < 1:
            nm = 32
        objs = arr[:, 4]
        cls_scores = arr[:, 5:5 + nc]
        scores_all = objs[:, None] * cls_scores
        best_cls = np.argmax(scores_all, axis=1)
        best_score = scores_all[np.arange(len(arr)), best_cls]
        mask = best_score >= conf
        if not np.any(mask):
            return []
        xywh = arr[mask, :4]
        boxes_raw = _xywh2xyxy(xywh)
        coeffs = arr[mask, 5 + nc:5 + nc + nm]
        scores = best_score[mask]
        clses = best_cls[mask]
    else:
        if arr.ndim == 3 and arr.shape[0] == 1:
            arr = arr[0]
        if arr.ndim != 2:
            return []
        if arr.shape[0] < arr.shape[1] and arr.shape[0] >= 4:
            arr = arr.T
        if arr.ndim != 2 or arr.shape[1] < 5:
            return []
        if nc is None:
            # 含 mask 系数：列数 = 4 + nc + nm，nm 一般 32
            nm = 32
            nc = arr.shape[1] - 4 - nm
        else:
            nm = arr.shape[1] - 4 - nc
        if nm < 1:
            nm = 32
            nc = arr.shape[1] - 4 - nm
        boxes_raw = _xywh2xyxy(arr[:, :4])
        cls_scores = arr[:, 4:4 + nc]
        coeffs = arr[:, 4 + nc:4 + nc + nm]
        best_cls = np.argmax(cls_scores, axis=1)
        best_score = cls_scores[np.arange(len(arr)), best_cls]
        mask = best_score >= conf
        if not np.any(mask):
            return []
        boxes_raw = boxes_raw[mask]
        coeffs = coeffs[mask]
        scores = best_score[mask]
        clses = best_cls[mask]

    boxes = [_scale_box(b, sx, sy, orig_w, orig_h) for b in boxes_raw]
    keep = _nms(boxes, scores, iou)
    out = []
    for i in keep:
        label = labels[int(clses[i])] if 0 <= int(clses[i]) < len(labels) else str(int(clses[i]))
        item = {"box": boxes[i], "label": str(label), "score": float(scores[i]), "task": "segment"}
        # mask 系数（简化：仅返回系数数组，由上层按需合成；不强制合成完整 mask 避免性能开销）
        try:
            item["mask_coeffs"] = [float(x) for x in coeffs[i]]
        except Exception:
            pass
        out.append(item)
    return out


# ===================== pose =====================
def _decode_pose(outputs, algorithm_type, labels, input_size, conf, iou, orig_size, nc=None):
    arr = np.asarray(outputs[0])
    iw, ih = input_size
    orig_w, orig_h = orig_size
    sx, sy = float(orig_w) / iw, float(orig_h) / ih
    # 仅 yolo8/11/26 pose 普遍为 [1, 4+1+nk*3, N]（4 box + 1 conf + nk*3 kpt）
    if arr.ndim == 3 and arr.shape[0] == 1:
        arr = arr[0]
    if arr.ndim != 2:
        return []
    if arr.shape[0] < arr.shape[1] and arr.shape[0] >= 5:
        arr = arr.T
    if arr.ndim != 2 or arr.shape[1] < 6:
        return []
    # 列布局：[cx, cy, w, h, conf, kpt_x, kpt_y, kpt_conf, ...] —— 标准 ultralytics pose 导出
    confs = arr[:, 4]
    mask = confs >= conf
    if not np.any(mask):
        return []
    arr = arr[mask]
    boxes_raw = _xywh2xyxy(arr[:, :4])
    confs = arr[:, 4]
    # 关键点：剩余列按 [x, y, conf] 三元组
    kpt_cols = arr.shape[1] - 5
    nk = kpt_cols // 3
    kpts = arr[:, 5:5 + nk * 3].reshape(-1, nk, 3) if nk > 0 else None
    boxes = [_scale_box(b, sx, sy, orig_w, orig_h) for b in boxes_raw]
    keep = _nms(boxes, confs, iou)
    out = []
    for i in keep:
        label = labels[0] if labels else "person"
        item = {"box": boxes[i], "label": str(label), "score": float(confs[i]), "task": "pose"}
        if kpts is not None:
            kp = kpts[i]
            # 缩放关键点 xy
            item["keypoints"] = [
                [float(kp[j, 0] * sx), float(kp[j, 1] * sy), float(kp[j, 2])]
                for j in range(nk)
            ]
        out.append(item)
    return out


# ===================== obb =====================
def _decode_obb(outputs, algorithm_type, labels, input_size, conf, iou, orig_size, nc=None):
    arr = np.asarray(outputs[0])
    iw, ih = input_size
    orig_w, orig_h = orig_size
    sx, sy = float(orig_w) / iw, float(orig_h) / ih
    # yolo8/11/26 obb: [1, 4+nc+1, N]  末列角度
    if arr.ndim == 3 and arr.shape[0] == 1:
        arr = arr[0]
    if arr.ndim != 2:
        return []
    if arr.shape[0] < arr.shape[1] and arr.shape[0] >= 5:
        arr = arr.T
    if arr.ndim != 2 or arr.shape[1] < 6:
        return []
    nc = arr.shape[1] - 5 if nc is None else nc
    if nc < 1:
        nc = 1
    boxes_raw = _xywh2xyxy(arr[:, :4])
    cls_scores = arr[:, 4:4 + nc]
    angles = arr[:, 4 + nc]
    best_cls = np.argmax(cls_scores, axis=1)
    best_score = cls_scores[np.arange(len(arr)), best_cls]
    mask = best_score >= conf
    if not np.any(mask):
        return []
    boxes_raw = boxes_raw[mask]
    scores = best_score[mask]
    clses = best_cls[mask]
    angles = angles[mask]
    boxes = [_scale_box(b, sx, sy, orig_w, orig_h) for b in boxes_raw]
    keep = _nms(boxes, scores, iou)
    out = []
    for i in keep:
        label = labels[int(clses[i])] if 0 <= int(clses[i]) < len(labels) else str(int(clses[i]))
        out.append({"box": boxes[i], "label": str(label), "score": float(scores[i]),
                    "task": "obb", "angle": float(angles[i])})
    return out