"""T9 SPIKE · BiRefNet ONNX 抠图管线 + 代理质检指标（T-P0-12 初步口径）。

代理指标（注意：这是代理口径，不是 T-P0-12 的正式验收——正式验收需要真实素材 +
人工标注的通过率，本 SPIKE 只建管线并给少量合成样图的主观核验依据）：
- 前景像素占比（合理区间 15%~60%）；
- 最大连通域占前景比（>0.9 视为「主体单一完整」）；
- 边缘羽化度（0.05<soft<0.95 过渡带面积 / 边界周长 ≈ 平均过渡带宽度 px）；
- 单张耗时（预热 3 + 正式 20 次；单张 >3s 时自动降为 10 次并注明）。

输入：合成人形图（纯色背景 + 渐变背景变体，不涉及真实个人数据）；
降级：models/ 下无抠图模型时用 Otsu 二值化占位（硬边缘，无羽化，报告需注明）；
输出：out/ 下掩膜图 + 抠图预览（gitignore，报告只留指标）+ matting_metrics.json。

用法：.venv/bin/python scripts/spikes/birefnet_matting.py
"""

from __future__ import annotations

import sys
import time
from typing import Any

import cv2
import numpy as np
import onnxruntime as ort
from synthetic import ASSETS_DIR, ensure_assets
from vision_common import MATTING_CHAIN, MODELS_DIR, OUT_DIR, bench, dump_json, load_session

# ImageNet 归一化（onnx-community 导出的 BiRefNet 系列 preprocessor_config 口径）
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

# 代理质检阈值（T-P2-07 质检雏形口径，非正式验收）
FG_RATIO_RANGE = (0.15, 0.60)  # 前景占比合理区间
LARGEST_CC_OK = 0.90  # 最大连通域占前景比阈值
SOFT_BAND = (0.05, 0.95)  # 软边过渡带判定区间
SLOW_RUNS = 10  # 单张 >3s 时的正式次数（默认 20）

_session: ort.InferenceSession | None = None
_model_source: str | None = None


def _get_session() -> ort.InferenceSession | None:
    """按候选链加载第一个可用模型（加载失败自动降级到下一候选）。"""
    global _session, _model_source
    if _model_source is not None:
        return _session
    for name in MATTING_CHAIN:
        path = MODELS_DIR / name
        if not path.exists():
            continue
        try:
            _session = load_session(path)
            _model_source = name
            print(f"[模型] 使用 {name}（{path.stat().st_size / 1048576:.1f} MB）")
            return _session
        except Exception as exc:  # SPIKE：逐候选降级，失败原因留痕后换下一个
            print(f"[降级] {name} 加载失败：{exc}")
    _session = None
    _model_source = "otsu-fallback"
    print("[降级] 无可用抠图模型，使用 Otsu 二值化占位（硬边缘、无羽化，报告需注明）")
    return None


def _preprocess(img_bgr: np.ndarray, size_wh: tuple[int, int]) -> np.ndarray:
    resized = cv2.resize(img_bgr, size_wh, interpolation=cv2.INTER_LINEAR)
    rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    rgb = (rgb - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(np.transpose(rgb, (2, 0, 1))[None])


def build_person_mask(img_bgr: np.ndarray) -> tuple[np.ndarray, str]:
    """抠图主入口：返回 (软掩膜 float32 HxW ∈ [0,1] 与原图同尺寸, 来源标签)。

    无模型时走 Otsu：人体比背景暗 → BINARY_INV 使人体=1。body_measure_prototype
    也复用本函数，保证两条 spike 的掩膜口径一致。
    """
    session = _get_session()
    if session is None:
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        thresh, _ = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        return thresh.astype(np.float32) / 255.0, "otsu-fallback"
    inp = session.get_inputs()[0]
    size_wh = (
        inp.shape[3] if isinstance(inp.shape[3], int) else 1024,
        inp.shape[2] if isinstance(inp.shape[2], int) else 1024,
    )
    tensor = _preprocess(img_bgr, size_wh)
    logits = session.run(None, {inp.name: tensor})[0]
    soft_small = 1.0 / (1.0 + np.exp(-np.clip(logits[0, 0].astype(np.float32), -30.0, 30.0)))
    soft = cv2.resize(
        soft_small, (img_bgr.shape[1], img_bgr.shape[0]), interpolation=cv2.INTER_LINEAR
    )
    assert _model_source is not None
    return soft, _model_source


def proxy_metrics(img_bgr: np.ndarray, soft: np.ndarray) -> dict[str, Any]:
    """T-P0-12 代理质检指标（合成图口径，正式验收需真实素材）。"""
    binary = (soft > 0.5).astype(np.uint8)
    fg = int(binary.sum())
    fg_ratio = fg / binary.size
    num_labels, _, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
    largest = int(stats[1:, cv2.CC_STAT_AREA].max()) if num_labels > 1 else 0
    largest_cc_ratio = largest / fg if fg else 0.0
    # 边缘羽化：过渡带面积 / 边界周长 ≈ 平均过渡带宽度（px）
    band = ((soft > SOFT_BAND[0]) & (soft < SOFT_BAND[1])).astype(np.uint8)
    band_area = int(band.sum())
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    perimeter = float(sum(cv2.arcLength(c, True) for c in contours)) or 1.0
    return {
        "fg_ratio": round(fg_ratio, 4),
        "fg_ratio_in_range": FG_RATIO_RANGE[0] <= fg_ratio <= FG_RATIO_RANGE[1],
        "largest_cc_ratio": round(largest_cc_ratio, 4),
        "largest_cc_ok": largest_cc_ratio > LARGEST_CC_OK,
        "feather_band_px": round(band_area / perimeter, 2),
        "band_area_px": band_area,
        "boundary_perimeter_px": round(perimeter, 1),
        "components": int(num_labels - 1),
    }


def _bench_latency(img_bgr: np.ndarray) -> dict[str, Any]:
    """单张全链路耗时（预处理+推理+sigmoid+resize 回原尺寸），口径同 rtmpose。"""

    def fn() -> tuple[np.ndarray, str]:
        return build_person_mask(img_bgr)

    fn()  # 预估单张耗时以自适应正式次数（BiRefNet 1024² 在 CPU 上可能较慢）
    t0 = time.perf_counter()
    fn()
    single = time.perf_counter() - t0
    runs = SLOW_RUNS if single > 3.0 else 20
    stats = bench(fn, warmup=3, runs=runs)
    stats["adaptive_note"] = f"单张约 {single:.2f}s → 正式 {runs} 次"
    return stats


def _save_outputs(stem: str, img: np.ndarray, soft: np.ndarray) -> None:
    """掩膜与抠图预览落盘（out/ 已 gitignore，仅作人工核验，不入库）。"""
    cv2.imwrite(str(OUT_DIR / f"mask_{stem}.png"), (soft * 255).astype(np.uint8))
    alpha = soft[..., None]
    cutout = (img.astype(np.float32) * alpha + 255.0 * (1.0 - alpha)).astype(np.uint8)
    cv2.imwrite(str(OUT_DIR / f"cutout_{stem}.png"), cutout)


def main() -> int:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    ensure_assets()
    images = {
        "synthetic_person": ASSETS_DIR / "synthetic_person.png",
        "synthetic_person_grad": ASSETS_DIR / "synthetic_person_grad.png",
    }
    results: dict[str, Any] = {}
    for stem, path in images.items():
        img = cv2.imread(str(path))
        soft, source = build_person_mask(img)
        metrics = proxy_metrics(img, soft)
        print(f"\n==== {stem}（{source}）====")
        print(
            f"  前景占比 {metrics['fg_ratio']:.4f}（区间 {FG_RATIO_RANGE}）"
            f" → {'OK' if metrics['fg_ratio_in_range'] else 'OUT'}"
        )
        print(
            f"  最大连通域/前景 {metrics['largest_cc_ratio']:.4f}（>{LARGEST_CC_OK}）"
            f" → {'OK' if metrics['largest_cc_ok'] else 'OUT'}；连通域数 {metrics['components']}"
        )
        print(f"  边缘羽化带 {metrics['feather_band_px']}px（过渡带 {metrics['band_area_px']}px²）")
        latency = _bench_latency(img)
        print(
            f"  耗时 runs={latency['runs']} 均值={latency['mean_s'] * 1000:.0f}ms "
            f"P95={latency['p95_s'] * 1000:.0f}ms 最大={latency['max_s'] * 1000:.0f}ms"
            f"（{latency['adaptive_note']}）"
        )
        _save_outputs(stem, img, soft)
        results[stem] = {"model_source": source, "metrics": metrics, "latency": latency}

    dump_json(OUT_DIR / "matting_metrics.json", results)
    print(f"\n结果已写入 {OUT_DIR / 'matting_metrics.json'}（掩膜/抠图预览在同目录）")
    return 0


if __name__ == "__main__":
    sys.exit(main())