"""DAgger 状态采集（2026-10-18 首次正式预训练后新增）：用当前策略在 train 任务上闭环跑（关安全层），
从它自己走到的状态里挑起点，交给老师重新规划 → 这些"策略会去、老师数据里没有"的状态拿到正确标签。
  python -m ml.dagger.collect --ckpt ml/runs/base/best.pt --samples ml/data/samples --n 8000 --seed 1 --parts 28 --out ml/data/dag1
每题最多取 2 个起点：
  - 撞了的题：首次碰撞前 0.2 s 的状态（策略"快要撞上"的地方，最有用）；
  - 另取一个 [0.3 s, 首撞前或结束] 内均匀随机时刻的状态（偏离老师的一般状态）。
起点必须无碰且间隙 > MIN_GAP_MM（老师从碰撞态出发规划不出来，与 build_dataset 的增广接受同一地板），且离目标 > 5 mm（否则 DLS 接管，不需要老师）。
输出 tasks_XX.jsonl（按 --parts 均分，供 gen_dataset.cjs --tasks 并行），每行 {q0, qPrev, T, R, cat, src}。
2026-10-19（审查 F7 / C10）：同时记录策略在该时刻的上一帧 qPrev；老师重规划的时间律从当前速度（v0）起步，
  t=0 样本的 dq 特征用真实上一帧——原来一律按静止起步标注，DAgger 样本不是策略真正到过的 (q, dq) 状态。
为什么需要：纯行为克隆在闭环里会偏出老师分布（首次 21850 步闭环到达 0%），原方案指望 APG/挖错修，但它们没法从 0% 起步。
老师只用于训练标注，部署仍然无老师。"""
import argparse, json, os, random
import torch

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

GAP_CHUNK = 1024      # 每次 kin.gaps 的查询点数上限


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True); ap.add_argument('--samples', nargs='+', default=['ml/data/samples'])
    ap.add_argument('--n', type=int, default=8000); ap.add_argument('--seed', type=int, default=1)
    ap.add_argument('--parts', type=int, default=8); ap.add_argument('--out', required=True)
    ap.add_argument('--bs', type=int, default=64)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    a = ap.parse_args()
    gap_min = float(os.environ.get('MIN_GAP_MM') or 3)
    os.makedirs(a.out, exist_ok=True)
    kin = Kin(device=a.device)
    from ml.common.ckpt import load_policy      # 2026-10-19：EMA 权重 + act_scale
    model, ck = load_policy(a.ckpt, a.device)
    assert ck.get('fp') in (None, kin.fp), '碰撞模型指纹与 checkpoint 不符'
    pool = [t for t in load_tasks(a.samples[:1]) if t['split'] == 'train' and t.get('cat') not in ('D',)]
    rng = random.Random(a.seed); rng.shuffle(pool); tasks = pool[:a.n]
    out, stat = [], {'tasks': len(tasks), 'collided': 0, 'states': 0, 'rej_gap': 0, 'rej_near': 0}
    for i in range(0, len(tasks), a.bs):
        tb = tasks[i:i + a.bs]
        with torch.no_grad():
            _, per = rollout(model, kin, tb, a.device, K=8, safety=False, seed=a.seed * 100003 + i, return_traj=True)
        tr = per['traj'].to(a.device)                                   # [B,T,6]
        B, Tn, _ = tr.shape
        _, T, R, _, _ = tasks_to_tensors(tb, a.device)
        # 2026-10-19 修（用户实测）：原来一次喂 64 题 × 64 步 = 4096 点、cap=40 mm 给 kin.gaps，40 mm 内幸存的胶囊对太多，
        # 峰值把显存 / 内存吃满。这里只需判断"≤0（碰撞）"和"> gap_min（起点地板）"，cap 取 gap_min+1 即可，
        # 且按 GAP_CHUNK 个点分块，用完即释放。
        cap = gap_min + 1.0
        flat = tr.reshape(-1, 6)
        with torch.no_grad():
            g = torch.cat([kin.gaps(flat[s:s + GAP_CHUNK], cap)['gm'].min(-1).values for s in range(0, flat.shape[0], GAP_CHUNK)]).view(B, Tn)
        del flat
        for b, t in enumerate(tb):
            hit = (g[b] <= 0).nonzero()
            c = int(hit[0]) if hit.numel() else Tn
            if hit.numel(): stat['collided'] += 1
            cand = []
            if hit.numel(): cand.append(max(0, c - int(round(0.2 / DT))))
            lo = int(round(0.3 / DT))
            if c - 1 > lo: cand.append(rng.randint(lo, c - 1))
            for k in cand:
                if g[b, k] <= gap_min: stat['rej_gap'] += 1; continue
                q = tr[b, k:k + 1]; qprev = tr[b, max(0, k - 1)]
                e, oe = goal_err(kin, q, T[b:b + 1], R[b:b + 1])
                if e.item() < 5 and (not t.get('R') or oe.item() < 5): stat['rej_near'] += 1; continue
                out.append({'q0': [round(v, 3) for v in q[0].tolist()], 'qPrev': [round(v, 3) for v in qprev.tolist()], 'T': t['T'], 'R': t.get('R'), 'cat': t.get('cat', 'A'),
                            'src': {'tid': t['tid'], 'k': k, 'col': bool(hit.numel())}})
        del tr, g, per
        if a.device.startswith('cuda'): torch.cuda.empty_cache()
        print(f'[dagger] {min(i + a.bs, len(tasks))}/{len(tasks)} 题，起点 {len(out)}', flush=True)
    stat['states'] = len(out)
    rng.shuffle(out)
    per_part = (len(out) + a.parts - 1) // max(1, a.parts)
    for p in range(a.parts):
        chunk = out[p * per_part:(p + 1) * per_part]
        if not chunk: continue
        fn = os.path.join(a.out, 'tasks_%02d.jsonl' % (p + 1))
        open(fn + '.tmp', 'w', encoding='utf-8').write(''.join(json.dumps(x) + '\n' for x in chunk)); os.replace(fn + '.tmp', fn)
    json.dump(stat, open(os.path.join(a.out, 'collect.json'), 'w', encoding='utf-8'), indent=1)   # 最后写：它存在 = 采集完成
    print(json.dumps(stat, ensure_ascii=False), flush=True)


if __name__ == '__main__':
    main()