"""T9 SPIKE · 共用小工具：基准统计、ORT 会话加载、RTMPose 推理解码封装。

只服务 scripts/spikes/ 下的四个脚本，不属于生产代码路径；网络访问仅存在于
download_models.py（本批唯一允许联网的环节），本模块与基准脚本全部离线可跑。
"""

from __future__ import annotations

import json
import subprocess
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import onnxruntime as ort

REPO_ROOT = Path(__file__).resolve().parents[2]
MODELS_DIR = REPO_ROOT / "models"
ASSETS_DIR = Path(__file__).resolve().parent / "assets"
OUT_DIR = Path(__file__).resolve().parent / "out"

# 模型候选链：与 download_models.py 的 TARGETS 保持一致；全缺失时调用方走降级（HOG/Otsu）
POSE_CHAIN: tuple[str, ...] = ("rtmpose-s.onnx", "rtmpose-m.onnx")
MATTING_CHAIN: tuple[str, ...] = ("birefnet_lite.onnx", "birefnet_lite_fp16.onnx", "rmbg14.onnx")

# RTMPose（mmpose 部署口径）：RGB 0~255 图像直接减均值除标准差
MMPOSE_MEAN = np.array([123.675, 116.28, 103.53], dtype=np.float32)
MMPOSE_STD = np.array([58.395, 57.12, 57.375], dtype=np.float32)

COCO17_NAMES = (
    "nose",
    "left_eye",
    "right_eye",
    "left_ear",
    "right_ear",
    "left_shoulder",
    "right_shoulder",
    "left_elbow",
    "right_elbow",
    "left_wrist",
    "right_wrist",
    "left_hip",
    "right_hip",
    "left_knee",
    "right_knee",
    "left_ankle",
    "right_ankle",
)
# COCO-17 索引（heights_from_keypoints / GT 校验用）
IDX_SHOULDER = (5, 6)
IDX_HIP = (11, 12)


def bench(fn: Callable[[], Any], warmup: int = 3, runs: int = 20) -> dict[str, Any]:
    """预热 warmup 次后正式 runs 次，返回均值/P95/单次最大（秒，perf_counter 口径）。

    T-P0-16 的验收口径就是这里的输出：预热 3 + 正式 20。
    """
    for _ in range(warmup):
        fn()
    samples: list[float] = []
    for _ in range(runs):
        t0 = time.perf_counter()
        fn()
        samples.append(time.perf_counter() - t0)
    arr = np.asarray(samples, dtype=np.float64)
    return {
        "runs": runs,
        "mean_s": round(float(arr.mean()), 4),
        "p95_s": round(float(np.percentile(arr, 95)), 4),
        "max_s": round(float(arr.max()), 4),
    }


def load_session(path: Path) -> ort.InferenceSession:
    """CPU EP 加载 ONNX。spike 约定不装 GPU 版 ORT（CUDA 匹配成本高，见任务文档 §2）。"""
    return ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])


def gpu_probe() -> dict[str, Any]:
    """探测 GPU 档可用性。任务文档：探测失败不折腾，只报 CPU。"""
    info: dict[str, Any] = {"ort_providers": list(ort.get_available_providers())}
    try:
        proc = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,driver_version", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=10,
            check=True,
        )
        info["nvidia_smi"] = proc.stdout.strip()
    except (OSError, subprocess.SubprocessError):
        info["nvidia_smi"] = None
    info["cuda_ep_available"] = "CUDAExecutionProvider" in info["ort_providers"]
    return info


def dump_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")


class RTMPoseRunner:
    """RTMPose ONNX（SimCC 输出）最小封装：整图 resize → 推理 → 关键点画布坐标。

    SPIKE 口径说明：自顶向下姿态模型本应由检测器给人体框再裁剪推理；合成图人形
    占满画布且位置已知，故直接整图缩放（等价于"完美检测框"），把变量集中在耗时
    测量上。输出 (K,3)：画布坐标 x、y + 置信度（x/y 两方向 softmax 峰值的较小者）。
    """

    def __init__(self, session: ort.InferenceSession) -> None:
        self.session = session
        inp = session.get_inputs()[0]
        self.input_name = inp.name
        # 动态维度（字符串）按 1 处理，仅用于确定输入尺寸
        shape = [d if isinstance(d, int) else 1 for d in inp.shape]
        if len(shape) != 4:
            msg = f"非预期的 RTMPose 输入维度: {inp.shape}"
            raise ValueError(msg)
        self.input_hw = (shape[2], shape[3])  # (H, W)

    def preprocess(self, img_bgr: np.ndarray) -> np.ndarray:
        h, w = self.input_hw
        resized = cv2.resize(img_bgr, (w, h), interpolation=cv2.INTER_LINEAR)
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32)
        rgb = (rgb - MMPOSE_MEAN) / MMPOSE_STD
        chw = np.transpose(rgb, (2, 0, 1))
        return np.ascontiguousarray(chw[None])

    def infer(self, img_bgr: np.ndarray) -> np.ndarray:
        """完整链路：预处理 + session.run + SimCC 解码 → (K,3)。"""
        tensor = self.preprocess(img_bgr)
        outputs = self.session.run(None, {self.input_name: tensor})
        return self._decode(outputs, orig_w=img_bgr.shape[1], orig_h=img_bgr.shape[0])

    def _decode(self, outputs: list[np.ndarray], orig_w: int, orig_h: int) -> np.ndarray:
        # SimCC 输出为两个 (1,K,Nx)/(1,K,Ny) 一维分布；Nx∝输入宽、Ny∝输入高
        cand = [o for o in outputs if o.ndim == 3]
        if len(cand) < 2:
            msg = f"非 SimCC 输出，形状: {[o.shape for o in outputs]}"
            raise ValueError(msg)
        in_h, in_w = self.input_hw
        a, b = cand[0], cand[1]
        # 与输入宽成比例的是 x 分布（比较两种配对的残差，取小者）
        fit_ab = abs(a.shape[2] / in_w - b.shape[2] / in_h)
        fit_ba = abs(b.shape[2] / in_w - a.shape[2] / in_h)
        if fit_ba < fit_ab:
            a, b = b, a
        x_simcc, y_simcc = a[0], b[0]  # (K,Nx) / (K,Ny)
        scale_x = a.shape[2] / in_w
        scale_y = b.shape[2] / in_h
        xs = x_simcc.argmax(axis=1).astype(np.float32) / scale_x
        ys = y_simcc.argmax(axis=1).astype(np.float32) / scale_y
        conf = np.minimum(_peak_prob(x_simcc), _peak_prob(y_simcc))
        kps = np.zeros((xs.shape[0], 3), dtype=np.float32)
        kps[:, 0] = xs * (orig_w / in_w)
        kps[:, 1] = ys * (orig_h / in_h)
        kps[:, 2] = conf
        return kps


def _peak_prob(simcc: np.ndarray) -> np.ndarray:
    """每个关键点一维分布的 softmax 峰值概率，作为置信度代理。"""
    e = np.exp(simcc - simcc.max(axis=1, keepdims=True))
    probs = e / e.sum(axis=1, keepdims=True)
    return probs.max(axis=1)