#!/usr/bin/env python
"""闭环末位误差诊断（只读，手动运行；不在流水线里）。

回答一个问题：策略"差一点没进 1 mm 门"还是"根本没朝目标走"。
输出每题：
  - 最终位置误差 err_mm / 最终姿态误差 oerr_deg
  - 全程最小位置误差 min_err_mm、最小综合误差时是否进过 5 mm & 5° 窗口（DLS 接管门槛）
  - done（是否到达）、碰撞步数、用时
以及按 cat/mode 分组的汇总与分位数。

用法（仓库根目录）：
  PYTHONUTF8=1 python ml/eval/diag_tail.py --ckpt ml/runs/base/last.pt --n 512 --out ml/runs/diag_tail.json
"""
import argparse, json, math, os, sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

import torch


def pct(v, p):
    if not v:
        return None
    v = sorted(v)
    return v[min(len(v) - 1, int(len(v) * p))]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True)
    ap.add_argument('--samples', default='ml/data/samples')
    ap.add_argument('--tasks', default=None, help='任务 jsonl（覆盖 --samples/--split）')
    ap.add_argument('--split', default='val')
    ap.add_argument('--n', type=int, default=512)
    ap.add_argument('--bs', type=int, default=64, help='每批题数（显存小就调小）')
    ap.add_argument('--K', type=int, default=8)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    ap.add_argument('--out', default=None)
    a = ap.parse_args()

    from ml.common.kin_torch import Kin, D2R
    from ml.data import load_tasks
    from ml.model.policy import build
    from ml.sim.closed_loop import rollout, goal_err

    from ml.common.ckpt import load_policy      # 2026-10-19：EMA 权重 + act_scale
    model, ck = load_policy(a.ckpt, a.device)
    kin = Kin(device=a.device)
    if ck.get('fp') and ck['fp'] != kin.fp:
        sys.exit('checkpoint 的碰撞模型指纹 %s 与当前 kin.json %s 不符' % (ck['fp'], kin.fp))

    if a.tasks:
        tasks = [json.loads(l) for l in open(a.tasks, encoding='utf-8') if l.strip()]
    else:
        tasks = [t for t in load_tasks([a.samples]) if t['split'] == a.split]
    tasks = tasks[:a.n]

    per_task = []
    with torch.no_grad():
        for i in range(0, len(tasks), a.bs):
            chunk = tasks[i:i + a.bs]
            res, per = rollout(model, kin, chunk, a.device, K=a.K, safety=False, seed=i, return_traj=True)
            traj = per['traj']                      # [B, T, 6]
            B, Tn, _ = traj.shape
            Tt = torch.tensor([t['T'] for t in chunk], dtype=torch.float32, device=a.device)
            Rt = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in chunk], dtype=torch.float32, device=a.device)
            min_e = torch.full((B,), float('inf'), device=a.device)
            min_o = torch.full((B,), float('inf'), device=a.device)
            entered = torch.zeros(B, dtype=torch.bool, device=a.device)   # 曾同时满足 e<5mm & (无姿态要求或 oe<5°)
            step = 256
            for t0 in range(0, Tn, step):
                q = traj[:, t0:t0 + step].reshape(-1, 6)
                n = q.shape[0] // B
                e, oe = goal_err(kin, q, Tt.repeat_interleave(n, 0), Rt.repeat_interleave(n, 0))
                e = e.view(B, n); oe = oe.view(B, n)
                min_e = torch.minimum(min_e, e.min(1).values)
                min_o = torch.minimum(min_o, oe.min(1).values)
                pose = (Rt.abs().sum(-1) > 0)[:, None]
                window = (e < 5) & (~pose | (oe < 5))
                entered |= window.any(1)
            done = per['done'].to(a.device); col = per['col'].to(a.device); td = per['t_done'].to(a.device)
            end_e = per['err'].to(a.device)
            # 末态姿态误差：traj 最后一帧
            qend = traj[:, -1]
            ee, eo = goal_err(kin, qend, Tt, Rt)
            for j, t in enumerate(chunk):
                per_task.append({
                    'tid': t.get('tid'), 'cat': t.get('cat'), 'mode': t.get('mode'),
                    'err_mm': float(end_e[j]), 'oerr_deg': float(eo[j]),
                    'min_err_mm': float(min_e[j]), 'min_oerr_deg': float(min_o[j]),
                    'entered_5mm5deg': bool(entered[j]),
                    'done': bool(done[j]), 'collision_steps': float(col[j]),
                    't_done_s': None if math.isnan(float(td[j])) else float(td[j]),
                })
            print('[diag] %d/%d  done=%d  entered=%d  min_err<=5mm=%d' % (
                min(i + a.bs, len(tasks)), len(tasks), int(done.sum()), int(entered.sum()),
                sum(1 for p in per_task if p['min_err_mm'] <= 5)), flush=True)

    def summ(rows):
        if not rows:
            return {'n': 0}
        e = [r['err_mm'] for r in rows]; me = [r['min_err_mm'] for r in rows]
        o = [r['oerr_deg'] for r in rows]
        return {'n': len(rows),
                'done': sum(r['done'] for r in rows),
                'entered_5mm5deg': sum(r['entered_5mm5deg'] for r in rows),
                'min_err_mm<=5': sum(1 for v in me if v <= 5),
                'min_err_mm<=10': sum(1 for v in me if v <= 10),
                'err_mm_p10/p50/p90': [pct(e, .1), pct(e, .5), pct(e, .9)],
                'min_err_mm_p10/p50/p90': [pct(me, .1), pct(me, .5), pct(me, .9)],
                'oerr_deg_p50/p90': [pct(o, .5), pct(o, .9)],
                'collision_steps_mean': sum(r['collision_steps'] for r in rows) / len(rows)}

    out = {'ckpt': a.ckpt, 'step': ck.get('step'), 'fp': ck.get('fp'), 'n': len(per_task),
           'overall': summ(per_task),
           'by_cat_mode': {k: summ([r for r in per_task if '%s/%s' % (r['cat'], r['mode']) == k])
                           for k in sorted({('%s/%s' % (r['cat'], r['mode'])) for r in per_task})},
           'tasks': per_task}
    s = json.dumps(out, indent=1, ensure_ascii=False)
    print(s[:2000])
    if a.out:
        open(a.out + '.tmp', 'w', encoding='utf-8').write(s)
        os.replace(a.out + '.tmp', a.out)
        print('[diag] 已写', a.out)


if __name__ == '__main__':
    main()