#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
8020 识别链路分段计时（249 上跑，不改生产代码，monkey-patch 打点）。

阶段划分（对应 query_all 一次调用）：
  read       读图 Image.open+np.array
  yolo       YOLO predict（含 letterbox/NMS，GPU）
  post       mask 二值化+质心+合并（detect_and_crop_all 总时 - yolo - crop）
  crop       逐实例 裁卡+warp（CPU）
  letterbox  裁卡图 letterbox392 + split_half（CPU，每卡）
  dino_u     上半区 DINOv2 前向→L2→cpu（GPU，每卡）
  dino_l     下半区 DINOv2 前向→L2→cpu（GPU，每卡）
  milvus_s   milvus upper_coll.search（网络召回K=30，每卡）
  milvus_q   milvus lower_coll.query（网络拉30条向量，每卡）
  rerank     下半区 numpy dot + fusion 排序（CPU，每卡；= milvus总 - s - q + rank）
  total      query_all 端到端

GPU 段天然有 .cpu().numpy() 同步点，perf_counter 直接可用。
每阶段记每次调用的 ms；输出 mean/p50 与占总时比例。

用法（pytorch env，249）:
  CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 \
    ~/miniconda3/envs/pytorch/bin/python tools/profile_8020_pipeline.py \
    --images data/audit_imgs/raspi_front_2026-09-01_13-22-48_529284.jpg \
             data/audit_imgs/multicard_2up.jpg --n 15
"""
import os
import sys
import time
import argparse
import statistics

sys.stdout.reconfigure(encoding="utf-8")
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, _ROOT)
sys.path.insert(0, os.path.join(_ROOT, "scripts"))

import numpy as np
from PIL import Image

import config
from scripts.cascade_match import DualCascadeMatcher
import scripts.cascade_match as cm

# 阶段 → [每次调用ms]
T = {}
def _rec(key, ms):
    T.setdefault(key, []).append(ms)


def _wrap(obj, attr, key, pre=None):
    """给 obj.attr 方法包一层计时；pre=objself 时传入宿主对象。"""
    fn = getattr(obj, attr)
    if getattr(fn, "_profile_wrapped", False):
        return

    def wrapper(*a, **kw):
        t0 = time.perf_counter()
        try:
            return fn(*a, **kw)
        finally:
            _rec(key, (time.perf_counter() - t0) * 1000)
    wrapper._profile_wrapped = True
    setattr(obj, attr, wrapper)


def install(matcher):
    # YOLO predict 全程
    _wrap(matcher.detector, "_predict", "yolo")
    # 裁卡+warp（每实例）
    _wrap(matcher.detector, "_crop_one_instance", "crop")
    # letterbox + split_half（cascade_match 命名空间里的引用，每卡一次）
    _wrap(cm, "letterbox392", "letterbox")
    _wrap(cm, "split_half", "split_half")
    # DINOv2 上/下区前向（每卡一次）
    _wrap(matcher.extractor, "extract_upper", "dino_u")
    _wrap(matcher.extractor, "extract_lower", "dino_l")
    # milvus 两次网络往返（每卡）——backend=milvus 时存在
    if hasattr(matcher, "upper_coll"):
        _wrap(matcher.upper_coll, "search", "milvus_s")
        _wrap(matcher.lower_coll, "query", "milvus_q")
    # gpu 常驻图库召回（backend=gpu 时存在）
    if hasattr(matcher, "gallery_upper"):
        _wrap(matcher, "_query_raw_gpu", "gpu_recall")
    # fusion 重排
    _wrap(matcher, "rank", "rank")


def run_once(matcher, path):
    t0 = time.perf_counter()
    matcher.query_all(path)
    total = (time.perf_counter() - t0) * 1000
    _rec("total", total)
    _rec("read", 0.0)  # 占位；read 实测并入下方独立测量


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--images", nargs="+", required=True)
    ap.add_argument("--n", type=int, default=15, help="正式计时次数（另有3次warmup）")
    ap.add_argument("--device", default="cuda:0")
    args = ap.parse_args()

    print(f"[Profile] 初始化 matcher（device={args.device}，与 8020 共存仅占显存不抢算力）...", flush=True)
    matcher = DualCascadeMatcher(
        alpha=config.CASCADE_ALPHA, top_k_recall=config.CASCADE_TOP_K_RECALL,
        device=__import__("torch").device(args.device),
        detector_device=__import__("torch").device(args.device),
        milvus_alias="prof", verbose=True,
    )
    install(matcher)

    # read 阶段单独测（query_all 内部的 Image.open 复刻）
    def measure_read(path, n):
        for _ in range(n):
            t0 = time.perf_counter()
            np.array(Image.open(path).convert("RGB"))
            _rec("read", (time.perf_counter() - t0) * 1000)

    ORDER = ["total", "read", "yolo", "post", "crop", "letterbox", "split_half",
             "dino_u", "dino_l", "milvus_s", "milvus_q", "gpu_recall", "rank"]

    for path in args.images:
        T.clear()
        for _ in range(3):
            run_once(matcher, path)
        T.clear()                       # warmup 丢弃
        for _ in range(args.n):
            run_once(matcher, path)
        measure_read(path, args.n)
        n_cards = len(matcher.query_all(path))  # 最终跑一次看卡数（计时已污染，仅取数）

        print(f"\n===== {os.path.basename(path)}  卡数={n_cards}  N={args.n} =====", flush=True)
        print(f"{'阶段':<12}{'mean ms':>9}{'p50 ms':>9}{'每次请求次数':>10}{'占总时':>8}", flush=True)
        total_mean = statistics.mean(T.get("total", [1]))
        rows = {}
        for k in ORDER:
            v = T.get(k)
            if not v:
                continue
            mean = statistics.mean(v)
            rows[k] = mean
            per_req = len(v) / max(1, args.n)
            print(f"{k:<12}{mean:>9.1f}{statistics.median(v):>9.1f}{per_req:>10.1f}"
                  f"{mean / total_mean:>7.1%}", flush=True)
        # post = detect_and_crop_all 整体由 yolo+crop+post 组成，post 无直接打点 → 用差值估
        # （detect_and_crop_all 未单独包计时，此处不拆，yolo/crop 已覆盖大头）

        gpu_sum = sum(rows.get(k, 0) for k in ("yolo", "dino_u", "dino_l", "gpu_recall"))
        cpu_sum = sum(rows.get(k, 0) for k in ("read", "post", "crop", "letterbox", "split_half", "rank"))
        io_sum = sum(rows.get(k, 0) for k in ("milvus_s", "milvus_q"))
        print(f"[归类] GPU(yolo+dino+召回)≈{gpu_sum:.0f}ms  CPU(读图/裁卡/letterbox/rank)≈{cpu_sum:.0f}ms  "
              f"Milvus IO≈{io_sum:.0f}ms  （每卡请求；总={total_mean:.0f}ms）", flush=True)

    print("\n[说明] post=mask二值化+质心+合并，未单独打点；total 与各段之和的差值即它+同步开销。"
          "\n[说明] dino_u/dino_l 为每卡一次：双卡图时每请求调 2 次，表内 mean 为单次调用耗时。", flush=True)


if __name__ == "__main__":
    main()