#!/usr/bin/env python
"""反馈蒸馏：用最终部署配置（目标反馈 + 安全层 + 部分前缀 + 每步规划）跑闭环，
收集 (特征, 复合动作 = 模型输出 + 目标反馈) 样本，供 train.py 微调把反馈内化进策略。
  PYTHONUTF8=1 python -m ml.dagger.distill_fb --ckpt ml/runs/base_dag2_dag1/best.pt --n 500 \
      --fb_gain 1.8 --out ml/data/samples_fbdistill
输出可直接被 ml.data.Samples 读取（collect.X/Y/V/SG/TID.npy；无 meta = act_scale 1.0，与主数据一致）。
动作标签 = 候选动作（含反馈）的归一化 act（/STEP，clamp ±1），与主数据同口径。
"""
import argparse, os, random
import numpy as np
import torch

from ml.common.ckpt import load_policy
from ml.common.features import STEP
from ml.common.kin_torch import Kin
from ml.data import load_tasks
from ml.sim.closed_loop import rollout


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True); ap.add_argument('--n', type=int, default=500)
    ap.add_argument('--bs', type=int, default=64); ap.add_argument('--fb_gain', type=float, default=1.8)
    ap.add_argument('--samples', default='ml/data/samples'); ap.add_argument('--out', required=True)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu'); ap.add_argument('--seed', type=int, default=1)
    a = ap.parse_args()
    dev = a.device
    model, _ = load_policy(a.ckpt, dev)
    kin = Kin(device=dev)
    pool = [t for t in load_tasks([a.samples]) if t['split'] == 'train' and t.get('cat') not in ('D',)]
    random.Random(a.seed).shuffle(pool); tasks = pool[:a.n]
    stp = torch.tensor(STEP)   # collect 返回 CPU 张量，标签换算也在 CPU
    X, Y, TID = [], [], []
    for i in range(0, len(tasks), a.bs):
        ch = tasks[i:i + a.bs]
        coll = []
        with torch.no_grad():
            rollout(model, kin, ch, dev, K=8, safety=True, replan='empty', fb_gain=a.fb_gain,
                    safety_partial=1, commit_n=1, seed=a.seed * 100003 + i, collect=coll)
        for f, ac in coll:
            X.append(f.numpy()); Y.append((ac / stp).clamp(-1, 1).numpy())
        ncol = sum(f.shape[0] for f, _ in coll)
        TID.append(np.arange(ncol, dtype=np.int32) + i * 100000)
        print(f'[distill] {i + len(ch)}/{len(tasks)} tasks, samples {sum(x.shape[0] for x in X)}', flush=True)
    os.makedirs(a.out, exist_ok=True)
    X = np.concatenate(X); Y = np.concatenate(Y); TID = np.concatenate(TID)
    n = X.shape[0]
    np.save(os.path.join(a.out, 'collect.X.npy'), X.astype(np.float32))
    np.save(os.path.join(a.out, 'collect.Y.npy'), Y.astype(np.float32))
    np.save(os.path.join(a.out, 'collect.V.npy'), np.full(n, 0.5, np.float32))
    np.save(os.path.join(a.out, 'collect.SG.npy'), np.zeros((n, 13), np.float32))
    np.save(os.path.join(a.out, 'collect.TID.npy'), TID)
    print(f'done: {n} samples, X {X.shape}, Y {Y.shape}', flush=True)


if __name__ == '__main__':
    main()