#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""_bench_index.py — 对比全量投影 locate() 与网格索引 TargetIndex 的耗时。"""
import time
import numpy as np

import target_math


def bench(fn, n):
    t0 = time.perf_counter()
    for _ in range(n):
        fn()
    return (time.perf_counter() - t0) / n * 1000.0


def main(path='acc.npz'):
    data = np.load(path)
    xyz = data['xyz']
    conf = data['confidence'] if 'confidence' in data else None
    valid = np.isfinite(xyz).all(axis=1) & (np.abs(xyz).sum(axis=1) > 1e-6)
    print('frame: %s  N=%d  valid=%d (%.1f%%)'
          % (path, xyz.shape[0], int(valid.sum()), 100.0 * valid.mean()))

    # 找一个确实能命中的像素
    query = None
    for u in range(700, 901, 20):
        for v in range(550, 751, 20):
            r = target_math.locate(xyz, u, v, conf=conf, radius=10)
            if r['found']:
                query = (float(u), float(v), r)
                break
        if query:
            break
    if query is None:
        raise SystemExit('没找到可命中的像素')
    u, v, r0 = query
    print('query pixel: (%.0f, %.0f)  p_lidar=%s  plane=%s  cand=%d'
          % (u, v, np.round(r0['p_lidar'], 4), r0['plane_used'], r0['n_cand']))

    n_click = 50
    t_full = bench(lambda: target_math.locate(xyz, u, v, conf=conf, radius=10), n_click)
    print('\n全量 locate():        %.3f ms/次（%d 次平均）' % (t_full, n_click))

    for cell in (16, 32, 64):
        t0 = time.perf_counter()
        idx = target_math.TargetIndex(xyz, conf=conf, cell=cell)
        build = (time.perf_counter() - t0) * 1000.0
        r1 = idx.locate(u, v, radius=10)
        t_idx = bench(lambda: idx.locate(u, v, radius=10), n_click)
        same = (r1['found'] == r0['found'])
        dp = None
        if r1['found'] and r0['found']:
            dp = np.linalg.norm(np.asarray(r1['p_lidar']) - np.asarray(r0['p_lidar']))
        print('cell=%2d px  建索引 %6.2f ms | 每次 %6.3f ms | 桶=%4d | 候选=%d | 结果一致=%s | Δp_lidar=%s'
              % (cell, build, t_idx, len(idx.buckets), r1['n_cand'], same,
                 '%.2e' % dp if dp is not None else '-'))

    # 模拟同一帧点 10 次（不同像素）
    rng = np.random.default_rng(0)
    pixels = [(u + rng.uniform(-60, 60), v + rng.uniform(-60, 60)) for _ in range(10)]

    def ten_full():
        for uu, vv in pixels:
            target_math.locate(xyz, uu, vv, conf=conf, radius=10)

    idx = target_math.TargetIndex(xyz, conf=conf, cell=32)
    idx.locate(u, v)  # 预热

    def ten_idx():
        for uu, vv in pixels:
            idx.locate(uu, vv, radius=10)

    t10_full = bench(ten_full, 10)
    t10_idx = bench(ten_idx, 10)
    print('\n同一帧点 10 次: 全量 %.2f ms | 索引 %.2f ms | 加速 %.1fx'
          % (t10_full, t10_idx, t10_full / max(t10_idx, 1e-9)))

    # 半径 20 的情况
    t20_idx = bench(lambda: idx.locate(u, v, radius=20), n_click)
    print('索引 radius=20:   %.3f ms/次' % t20_idx)


if __name__ == '__main__':
    import sys
    main(sys.argv[1] if len(sys.argv) > 1 else 'acc.npz')