"""老师轨迹回放自检（train 阶段开训前自动跑，约 1–3 分钟）：
  python -m ml.eval.replay_teacher --n 512 --out ml/runs/replay_teacher.json
两部分，都过才 pass（pipeline.sh 读 pass 字段与指纹，不过就停 rc 5）：
  1. 运动学口径（2026-10-18）：gen 写出的老师 50 Hz 轨迹（keep=true、cat≠D）直接放进 PyTorch 运动学 / 碰撞 / 到达判据：
     到达率 ≥ 0.95、碰撞率 ≤ 0.01。验证的是「仿真口径（kin.json、间隙、到达门槛）与 JS 认证过的数据一致」。
  2. 滚动口径（2026-10-19，审查 review_closedloop F1/D1、review_data_actions C1）：上面那条完全不经过 rollout()，
     队列 / 前缀 / 提交 3 步 / qe / 标签 clamp / DLS / done / 时限上的错位它都发现不了。这里让老师当策略（closed_loop.TeacherOracle：
     每次规划从当前时间下标取老师 16 步增量，按标签口径 clamp）原样走 rollout()（select='first'、use_sub=False）：
     到达（时限内）≥ 0.95、碰撞 ≤ 0.01、DLS 接管前与老师轨迹的最大关节偏差 p99 ≤ 0.1°（shard 轨迹按 1e-3° 取整，标签 clamp 会截掉取整的余量）。
输出另带 fp（当前 kin.json）、data_fp（shard 里记录的碰撞模型指纹）、ver=2；pipeline 三者都核对。
load_teacher / pad_traj / oracle_raw / oracle_summary 供 ml/eval/diag_cl.py 的 H1 段复用。"""
import argparse, glob, json, math, os, sys
import torch

from ml.common.kin_torch import Kin
from ml.sim.closed_loop import goal_err, rollout, TeacherOracle, DT

DEV_TOL_DEG = 0.1


def load_teacher(shards, ids):
    """按任务 id 从 shard jsonl 取老师记录（traj、path、exec_s…）。只解析 id 命中的行。"""
    want, out = set(ids), {}
    for f in sorted(glob.glob(shards)):
        for line in open(f, encoding='utf-8'):
            k = line.find('"id":"')
            if k < 0: continue
            rid = line[k + 6:line.find('"', k + 6)]
            if rid in want and rid not in out:
                r = json.loads(line)
                if r.get('traj'): out[rid] = r
        if len(out) == len(want): break
    return out


def pad_traj(trajs, device, extra=20):
    """[B,Tmax+extra,6]：各题老师轨迹末帧重复补齐（oracle 在轨迹末尾之后给零增量）。"""
    Tm = max(len(t) for t in trajs) + extra
    out = torch.zeros(len(trajs), Tm, 6, device=device)
    for i, t in enumerate(trajs):
        x = torch.tensor(t, dtype=torch.float32, device=device)
        out[i, :len(t)] = x; out[i, len(t):] = x[-1]
    return out


def oracle_raw(per, TR, lens, tasks):
    """oracle 滚动的逐题量：DLS 接管前（且到达前、老师轨迹结束前）与老师的最大关节偏差 mdev（度）、时限内到达 ok、
    时限内碰撞 col、是否计入到达统计（cat≠D）、完成时刻 − 老师时长 dt（s）、是否用过 DLS、2 s 目标误差比。"""
    tr = per['traj'].to(TR.device); B, Tn, _ = tr.shape
    L = min(Tn, TR.shape[1])
    dev = (tr[:, :L] - TR[:, :L]).abs().max(-1).values                         # [B,L]
    k = torch.arange(L, device=TR.device)[None]
    kd = torch.tensor(lens, device=TR.device)[:, None] - 1
    dls = per['dls_it'].to(TR.device)[:, None]
    td = per['t_done'].to(TR.device)
    kdone = torch.where(torch.isfinite(td), (td / DT).round(), torch.full_like(td, 1e9)).long()[:, None]
    valid = (k <= kd) & ((dls < 0) | (k <= dls)) & (k <= kdone)
    ex = torch.tensor([t.get('exec_s') or float('nan') for t in tasks])
    return {'mdev': torch.where(valid, dev, torch.zeros_like(dev)).max(-1).values.cpu(), 'ok': per['ok'], 'col': per['col'] > 0,
            'reach_m': torch.tensor([t.get('cat') != 'D' for t in tasks]), 'dt': per['t_done'] - ex,
            'dls': per['dls_it'] >= 0, 'r2': per['goal_ratio_2s']}


def _q(v, p):
    v = v[torch.isfinite(v)]
    return float(v.kthvalue(max(1, min(len(v), int(math.ceil(p * len(v)))))).values) if len(v) else None


def oracle_summary(raws):
    c = {k: torch.cat([r[k] for r in raws]) for k in raws[0]}
    m = c['reach_m']; ok = c['ok']
    dt = c['dt'][ok & m]
    return {'n': int(len(m)), 'n_reach_tasks': int(m.sum()),
            'reach': float(ok[m].float().mean()) if m.any() else None,
            'collision': float(c['col'][m].float().mean()) if m.any() else None,
            'D': {'n': int((~m).sum()), 'collision': float(c['col'][~m].float().mean()) if (~m).any() else None},
            'max_dev_deg_before_dls': {'p50': _q(c['mdev'], .5), 'p99': _q(c['mdev'], .99), 'max': float(c['mdev'].max())},
            'dls_used': float(c['dls'].float().mean()),
            'finish_minus_teacher_s': {'p50': _q(dt, .5), 'p90': _q(dt, .9), 'max': float(dt.max()) if len(dt) else None},
            'goal_ratio_2s_p50': _q(c['r2'], .5)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--shards', default='ml/data/shard_*.jsonl'); ap.add_argument('--n', type=int, default=512)
    ap.add_argument('--samples', default='ml/data/samples', help='读打包 meta 的 act_scale（oracle 按同一标签口径 clamp）；没有就按 1.0')
    ap.add_argument('--bs', type=int, default=128)
    ap.add_argument('--out', default='ml/runs/replay_teacher.json')
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    a = ap.parse_args()
    kin = Kin(device=a.device)
    recs = []
    for f in sorted(glob.glob(a.shards)):
        for line in open(f, encoding='utf-8'):
            r = json.loads(line)
            if r.get('keep') and r.get('traj') and r.get('cat') != 'D': recs.append(r)
            if len(recs) >= a.n: break
        if len(recs) >= a.n: break
    if not recs: sys.exit('没有可回放的老师轨迹（先跑 gen）')
    # ---- 1. 运动学口径
    reach = col = 0; worst_gap = []
    for r in recs:
        q = torch.tensor(r['traj'], dtype=torch.float32, device=a.device)
        mid = 0.5 * (q[1:] + q[:-1])
        g = kin.gaps(torch.cat([q, mid], 0), 0.5)['gm'].min(-1).values
        hit = bool((g <= 0).any()); col += hit; worst_gap.append(g.min().item())
        T = torch.tensor([r['T']], dtype=torch.float32, device=a.device)
        R = torch.tensor([r['R'] if r.get('R') else [0.] * 9], dtype=torch.float32, device=a.device)
        e, oe = goal_err(kin, q[-1:], T, R)
        reach += bool((e < 1).item() and (oe < 0.5).item())
    n = len(recs)
    kin_ok = reach / n >= 0.95 and col / n <= 0.01
    # ---- 2. 滚动口径：老师 oracle 走完整 rollout()
    try:
        from ml.data import read_meta
        act_scale = float(read_meta(a.samples).get('act_scale', 1.0)) if os.path.isdir(a.samples) else 1.0
    except SystemExit:
        act_scale = 1.0
    tasks = [{'q0': r['q0'], 'T': r['T'], 'R': r.get('R'), 'exec_s': r.get('exec_s') or (len(r['traj']) - 1) * DT,
              'travel': r.get('travel_deg'), 'cat': r.get('cat')} for r in recs]
    raws = []
    for i in range(0, n, a.bs):
        tb, rb = tasks[i:i + a.bs], recs[i:i + a.bs]
        TR = pad_traj([r['traj'] for r in rb], a.device)
        _, per = rollout(None, kin, tb, a.device, K=1, select='first', use_sub=False, return_traj=True, seed=i,
                         oracle=TeacherOracle(TR, act_scale=act_scale))
        raws.append(oracle_raw(per, TR, [len(r['traj']) for r in rb], tb))
    orc = oracle_summary(raws); orc['act_scale'] = act_scale
    orc_ok = (orc['reach'] or 0) >= 0.95 and (orc['collision'] or 0) <= 0.01 and orc['max_dev_deg_before_dls']['p99'] <= DEV_TOL_DEG
    data_fps = {r.get('fp') for r in recs}
    data_fp = data_fps.pop() if len(data_fps) == 1 else 'mixed'
    res = {'ver': 2, 'n': n, 'reach': reach / n, 'collision': col / n, 'min_gap_mm_p10': sorted(worst_gap)[n // 10], 'kin_pass': kin_ok,
           'oracle': orc, 'oracle_pass': orc_ok, 'fp': kin.fp, 'data_fp': data_fp,
           'pass': bool(kin_ok and orc_ok and data_fp in (None, kin.fp))}
    os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
    json.dump(res, open(a.out, 'w', encoding='utf-8'), indent=1, ensure_ascii=False)
    print('老师轨迹回放：到达 %.1f%%、碰撞 %.1f%%（%d 条）；oracle 走 rollout：到达 %.1f%%、碰撞 %.1f%%、DLS 前最大偏差 p99 %.4f°；指纹 %s' % (
        100 * res['reach'], 100 * res['collision'], n, 100 * orc['reach'], 100 * orc['collision'],
        orc['max_dev_deg_before_dls']['p99'], '一致' if data_fp in (None, kin.fp) else '不符'), flush=True)
    sys.exit(0 if res['pass'] else 5)


if __name__ == '__main__':
    main()