/* 原始老师轨迹 → 闭环训练样本（方案 §4.5 / §5）
 *   node ml/prep/build_dataset.cjs ml/data/shard_0001.jsonl [outDir=ml/data/samples]
 * 每条 keep=true 的轨迹按 50 Hz 时间律重放，每 3 步取一个样本：
 *   X  [N,83]   特征（webapp/js/ik-policy.js features()，与部署同一份代码）
 *   Y  [N,16,6] 未来 16 步每步增量 / STEP（到终点后补 0）
 *   V  [N]      剩余时间（s）/ 5
 *   SG [N,13]   子目标标签：剩余两个途经点相对当前 q 的差 / 180（不足用终点补）+ 是否需要
 *   TID[N]      任务编号（tasks.jsonl 行号所在的全局 id，用于划分与闭环验证）
 * 增广（2026-10-18 首次正式预训练后加大）：40% 样本对 q_t 加 σ=AUG_SIGMA（默认 3°）噪声，标签改为 AUG_BACK（默认 12）步内线性收回老师轨迹；
 *   噪声态须无干涉且间隙 > MIN_GAP_MM，收回途中第 3/6/9/12 步也须无干涉（σ 加大后直接线性拉回可能穿过障碍）。
 *   原为 30%、σ=1°、8 步：闭环一偏出 1° 就是没见过的状态，首次 21850 步闭环到达 0%。更大的偏离由 DAgger 阶段补（ml/dagger）。
 *       已提交前缀长度 r ∈ {0..3} 均匀随机，前缀取老师接下来 r 步的增量。 */
'use strict';
const fs = require('fs'), path = require('path');
const IN = process.argv[2], OUTDIR = process.argv[3] || path.join(__dirname, '../data/samples');
if (!IN) { console.error('usage: build_dataset.cjs <shard.jsonl> [outDir]'); process.exit(1); }
const { S } = require('../../tests/route_bench_setup.cjs')();
const POL = require('../../webapp/js/ik-policy.js');
const { timeParam } = require('../gen/timelaw.cjs');
const { writeNpy } = require('./npy.cjs');
const { F_DIM, STEP, H } = POL;
const L = S.LIM, O = S.ORDER;
const base = path.basename(IN, '.jsonl');
let st = 0; for (const ch of base) st = (st * 31 + ch.charCodeAt(0)) >>> 0;
const rnd = () => { st = (st * 1664525 + 1013904223) >>> 0; return st / 4294967296; };
const gauss = () => Math.sqrt(-2 * Math.log(rnd() + 1e-12)) * Math.cos(2 * Math.PI * rnd());
const hash = s => { let h = 2166136261; for (const c of s) h = Math.imul(h ^ c.charCodeAt(0), 16777619); return h >>> 0; };
const X = [], Y = [], V = [], SG = [], TID = [], tasks = [];
const lines = fs.readFileSync(IN, 'utf8').split('\n').filter(Boolean);
let nAug = 0, nAugRej = 0, nGapDrop = 0;
/* 2026-10-18：认证最小间隙 < 3 mm 的老师轨迹不进训练集；噪声增广接受地板同步为 3 mm（原 1 mm）。 */
const MIN_GAP_MM = +(process.env.MIN_GAP_MM || 3);
const AUG_P = +(process.env.AUG_P || 0.4), AUG_SIGMA = +(process.env.AUG_SIGMA || 3), AUG_BACK = +(process.env.AUG_BACK || 12);
const NT = lines.length, PSTEP = Math.max(1, Math.floor(NT / 10));
for (let li = 0; li < NT; li++) {
  if (li % PSTEP === 0) console.log(`[prep] ${base} ${li}/${NT}`);   // 每 10% 一行，28 个分片并行也看得见进度
  const line = lines[li];
  const r = JSON.parse(line);
  if (!r.keep) continue;
  if (r.cert && r.cert.minGap != null && r.cert.minGap < MIN_GAP_MM) { nGapDrop++; continue; }   // 与 stats.cjs 同一过滤
  const tid = hash(r.id) & 0x7fffffff;
  const split = (hash(r.id + '#split') % 100) < 5 ? 'val' : 'train';
  const P = r.path;
  const tp = timeParam(P);
  const traj = tp.traj, taus = tp.taus, N = traj.length - 1;
  tasks.push({ tid, id: r.id, cat: r.cat, mode: r.mode, q0: r.q0, T: r.T, R: r.R, exec_s: N * POL.DT, travel: r.travel_deg, reached: r.reached, err: r.err, split });
  const vt = tp.cum;                                   // 顶点 τ
  const inter = P.slice(1, -1).map((p, i) => ({ p, tau: vt[i + 1] }));
  for (let t = 0; t <= N; t += 3) {
    let q = traj[t].slice(), qPrev = (t > 0 ? traj[t - 1] : traj[0]).slice();
    let off = [0, 0, 0, 0, 0, 0], aug = false;
    if (rnd() < AUG_P) {
      const qn = q.map((v, j) => { let x = v + AUG_SIGMA * gauss(); const t2 = L[O[j]]; if (!t2.cont) x = Math.max(t2.lo, Math.min(t2.hi, x)); return x; });
      nAug++;
      const o2 = qn.map((v, j) => v - q[j]);
      const back = k => traj[Math.min(N, t + k)].map((v, j) => v + o2[j] * Math.max(0, 1 - k / AUG_BACK));
      if (S.collideFree(qn) && S.clearance(qn, 5) > MIN_GAP_MM && [3, 6, 9, 12].every(k => k >= AUG_BACK || S.collideFree(back(k)))) {
        off = o2; q = qn; qPrev = qPrev.map((v, j) => v + off[j]); aug = true;
      } else nAugRej++;
    }
    const pos = k => { const i = Math.min(N, t + k); const w = aug ? Math.max(0, 1 - k / AUG_BACK) : 0; return traj[i].map((v, j) => v + off[j] * w); };
    const y = new Float32Array(H * 6);
    let prevP = q;
    for (let k = 1; k <= H; k++) { const pk = pos(k); for (let j = 0; j < 6; j++) y[(k - 1) * 6 + j] = Math.max(-1, Math.min(1, (pk[j] - prevP[j]) / STEP[j])); prevP = pk; }
    const rr = Math.floor(rnd() * 4), prefix = [];
    for (let k = 0; k < rr; k++) { const row = []; for (let j = 0; j < 6; j++) row.push(y[k * 6 + j] * STEP[j]); prefix.push(row); }
    const remain = inter.filter(v => v.tau > taus[t] + 0.02);
    const sub = remain.length && rnd() < 0.8 ? remain[0].p : null;
    const gm = S.gapMatrix(q, POL.GAP_CAP);
    const f = POL.features(S, { q, qPrev, T: r.T, R: r.R, prefix, sub }, gm);
    const v1 = remain[0] ? remain[0].p : P[P.length - 1], v2 = remain[1] ? remain[1].p : P[P.length - 1];
    const sg = new Float32Array(13);
    for (let j = 0; j < 6; j++) { sg[j] = (v1[j] - q[j]) / 180; sg[6 + j] = (v2[j] - q[j]) / 180; }
    sg[12] = remain.length ? 1 : 0;
    X.push(f); Y.push(y); V.push((N - t) * POL.DT / 5); SG.push(sg); TID.push(tid);
  }
}
fs.mkdirSync(OUTDIR, { recursive: true });
const n = X.length;
const cat = (arrs, w) => { const o = new Float32Array(n * w); arrs.forEach((a, i) => o.set(a, i * w)); return o; };
writeNpy(path.join(OUTDIR, base + '.X.npy'), cat(X, F_DIM), [n, F_DIM]);
writeNpy(path.join(OUTDIR, base + '.Y.npy'), cat(Y, H * 6), [n, H, 6]);
writeNpy(path.join(OUTDIR, base + '.V.npy'), Float32Array.from(V), [n]);
writeNpy(path.join(OUTDIR, base + '.SG.npy'), cat(SG, 13), [n, 13]);
writeNpy(path.join(OUTDIR, base + '.TID.npy'), TID, [n], 'int32');
// 2026-10-18 断点续跑：tasks.jsonl 最后写、先写 .tmp 再改名 —— 它存在就代表这个 shard 的 5 个 npy 都已写完，run_prep.sh 据此跳过
const tf = path.join(OUTDIR, base + '.tasks.jsonl');
fs.writeFileSync(tf + '.tmp', tasks.map(t => JSON.stringify(t)).join('\n') + '\n'); fs.renameSync(tf + '.tmp', tf);
console.log(JSON.stringify({ shard: base, tasks: tasks.length, samples: n, aug: nAug, augRejected: nAugRej, gapDrop: nGapDrop, featVer: POL.FEAT_VER }));