#!/usr/bin/env node
// Refit the difficulty weights that select-50 scores with, from a round of
// agent play.
//
//   node scripts/refit-weights.mjs [--agent="R1 · grok xhigh"] [--ridge=0.3]
//
// select-50 scores a candidate from mechanical features, but the weights on
// those features started as guesses. This fits them against what agents
// actually experienced, and prints the --weights flag for the next round.
//
// The response is not the difficulty rating alone. That rating is reliable
// (r = 0.93 on retest) but it is a 1-5 integer, so on fifty levels it is nearly
// constant and a fit against it alone would be fitting rounding. It is averaged
// with two continuous measures of the same thing — how many decisions the level
// forced, and how much action went nowhere — each standardised first.
//
// Ridge, because seven features on fifty levels with collinear inputs (bigger
// boards hold more pieces) will otherwise hand back large coefficients that
// cancel. The printed R² says whether the fit earned its keep; when it is near
// zero, the honest move is to keep the old weights and say the round carried no
// signal.
import { existsSync, readFileSync, readdirSync } from 'node:fs';
import { resolve } from 'node:path';
import { metrics } from '../src/bench.js';
import { legibility } from '../src/deep-lock.js';

const args = Object.fromEntries(process.argv.slice(2).map(a => { const [k, ...v] = a.replace(/^--/, '').split('='); return [k, v.length ? v.join('=') : true]; }));
const AGENT = args.agent ?? 'R1 · grok xhigh';
const RIDGE = Number(args.ridge ?? 0.3);

// The features select-50 scores with, plus the two the design corrections
// argued matter — how much the opening offers, and how much a winding route
// pays off over a straight one.
const FEATURES = ['rotations', 'pieces', 'lockedShare', 'monkeyAPT', 'cells', 'contrast', 'openingChoices',
  'glanceCost', 'medianVerify', 'grazeTraps'];
const featureOf = level => {
  const c = level.campaign ?? {};
  const p = level.prototype ?? {};
  return {
    rotations: c.rotations ?? 3, pieces: c.pieces ?? level.tiles.length,
    lockedShare: c.lockedShare ?? 0, monkeyAPT: p.monkeyActionsPerPiece ?? 0,
    cells: level.cells.length, contrast: c.contrast ?? 0,
    openingChoices: c.openingChoices ?? 0,
    ...(({ glanceCost, medianVerify, grazeTraps }) => ({ glanceCost: glanceCost ?? 0, medianVerify, grazeTraps }))(legibility(level)),
  };
};

const levels = JSON.parse(readFileSync(resolve('levels/campaign-50.json'), 'utf8')).levels;
const byId = new Map(levels.map(l => [l.id, l]));

const played = new Map();
const RUNS = resolve('bench/runs');
for (const f of existsSync(RUNS) ? readdirSync(RUNS) : []) {
  if (!f.endsWith('.json') || f === 'index.json') continue;
  try {
    const d = JSON.parse(readFileSync(resolve(RUNS, f), 'utf8'));
    if (d.agent !== AGENT || !byId.has(d.level?.id)) continue;
    played.set(d.level.id, metrics(d));
  } catch { /* torn read */ }
}

const rows = [...played].map(([id, m]) => ({ id, m, x: featureOf(byId.get(id)) }))
  .filter(r => Number.isFinite(Number(r.m.rating?.difficulty)));
if (rows.length < 12) {
  console.error(`只有 ${rows.length} 关有评分，样本太少，不拟合。`);
  process.exit(1);
}

const mean = a => a.reduce((x, y) => x + y, 0) / a.length;
const sd = a => { const m = mean(a); return Math.sqrt(mean(a.map(v => (v - m) ** 2))); };
const z = a => { const m = mean(a), s = sd(a); return s === 0 ? a.map(() => 0) : a.map(v => (v - m) / s); };

// Three views of the same quantity: what the agent said, how much it had to
// decide, and how much of what it did led nowhere.
const parts = {
  difficulty: rows.map(r => Number(r.m.rating.difficulty)),
  decisions: rows.map(r => Math.log1p(r.m.decisions ?? 0)),
  wasted: rows.map(r => Math.log1p((r.m.deadTaps ?? 0) + (r.m.blindRotations ?? 0) + (r.m.wrongWayRotations ?? 0))),
};
// A part with no spread contributes nothing but still divides the average, so
// it is dropped by name rather than silently zeroed. This matters: a round
// where every level came back rated 2 leaves an effort-only response, and
// effort alone largely re-measures board size.
const liveParts = Object.entries(parts).filter(([, v]) => sd(v) > 1e-9);
const flatParts = Object.keys(parts).filter(k => !liveParts.some(([n]) => n === k));
if (!liveParts.length) { console.error('本轮没有任何一项有方差，无从拟合。'); process.exit(1); }
const zparts = liveParts.map(([, v]) => z(v));
const y = rows.map((_, i) => mean(zparts.map(p => p[i])));
const zy = z(y);

// Drop features the round could not separate: a column that never varies
// carries no information and a ridge fit would still hand it a coefficient.
const live = FEATURES.filter(f => sd(rows.map(r => r.x[f])) > 1e-9);
const dropped = FEATURES.filter(f => !live.includes(f));
const cols = live.map(f => rows.map(r => r.x[f]));
const sds = Object.fromEntries(live.map((f, j) => [f, sd(cols[j])]));
const Z = live.map((_, j) => z(cols[j]));

// beta = (Z'Z + λI)^-1 Z'y, solved by Gaussian elimination with partial pivoting.
const p = live.length, n = rows.length;
const A = Array.from({ length: p }, (_, i) => Array.from({ length: p + 1 }, (_, j) => {
  if (j === p) return Z[i].reduce((s, v, k) => s + v * zy[k], 0) / n;
  return Z[i].reduce((s, v, k) => s + v * Z[j][k], 0) / n + (i === j ? RIDGE : 0);
}));
for (let i = 0; i < p; i++) {
  let piv = i;
  for (let k = i + 1; k < p; k++) if (Math.abs(A[k][i]) > Math.abs(A[piv][i])) piv = k;
  [A[i], A[piv]] = [A[piv], A[i]];
  for (let k = 0; k < p; k++) {
    if (k === i || A[i][i] === 0) continue;
    const f = A[k][i] / A[i][i];
    for (let j = i; j <= p; j++) A[k][j] -= f * A[i][j];
  }
}
const beta = A.map((row, i) => row[i] === 0 ? 0 : row[p] / row[i]);

const fitted = rows.map((_, k) => live.reduce((s, _f, j) => s + beta[j] * Z[j][k], 0));
const ssRes = zy.reduce((s, v, k) => s + (v - fitted[k]) ** 2, 0);
const r2 = 1 - ssRes / zy.reduce((s, v) => s + v * v, 0);

// Back to raw-feature units, then rescaled so scores land in the same range as
// the round just played — the score is only ever used to rank and to spread
// picks evenly, so its scale is free, and holding it steady keeps the printed
// band numbers comparable from round to round.
const OLD = { rotations: 1.0, pieces: 0.10, lockedShare: 1.6, monkeyAPT: 0.14, cells: 0.004, contrast: 0, openingChoices: 0, glanceCost: 0, medianVerify: 0, grazeTraps: 0 };
const raw = Object.fromEntries(live.map((f, j) => [f, beta[j] / sds[f]]));
const scoreWith = w => rows.map(r => live.reduce((s, f) => s + (w[f] ?? 0) * r.x[f], 0));
const gain = sd(scoreWith(OLD)) / (sd(scoreWith(raw)) || 1);
const fit = Object.fromEntries(live.map(f => [f, Number((raw[f] * gain).toFixed(4))]));
for (const f of dropped) fit[f] = OLD[f] ?? 0;

const corr = (a, b) => { const za = z(a), zb = z(b); return mean(za.map((v, i) => v * zb[i])); };
console.log(`样本 ${n} 关 · agent ${AGENT} · ridge ${RIDGE}`);
const NAME = { difficulty: '评分难度', decisions: 'log 决策数', wasted: 'log 无效动作' };
console.log(`合成难度 = ${liveParts.map(([k]) => 'z(' + NAME[k] + ')').join(' + ')} 的均值`);
if (flatParts.length) {
  console.log(`  !! ${flatParts.map(k => NAME[k]).join('、')} 本轮方差为 0，没参与合成`);
  if (flatParts.includes('difficulty')) {
    console.log('  !! 难度评分全同分 —— 这一轮拿不到难度信号，剩下的只是「这关有多少事要做」，');
    console.log('     而那主要由棋子数和格子数决定。照这组权重重选，等于按盘面大小排序。');
  }
}
for (const [k, v] of liveParts) for (const [k2, v2] of liveParts) if (k < k2) console.log(`  ${NAME[k]} × ${NAME[k2]}：${corr(v, v2).toFixed(2)}`);
console.log(`拟合 R² = ${r2.toFixed(3)}${r2 < 0.25 ? '  ← 太低，这一轮没给出可用信号，别换权重' : ''}`);
if (flatParts.includes('difficulty')) console.log('R² 高不代表可用：被拟合的是工作量，不是难度。');
console.log('');
console.log('| 特征 | 单独相关 | 标准化系数 | 旧权重 | 新权重 |');
console.log('|---|---:|---:|---:|---:|');
for (const f of FEATURES) {
  const j = live.indexOf(f);
  const c = j < 0 ? null : corr(rows.map(r => r.x[f]), y);
  console.log(`| ${f} | ${c == null ? '—（本轮无变化）' : c.toFixed(2)} | ${j < 0 ? '—' : beta[j].toFixed(3)} | ${(OLD[f] ?? 0)} | ${fit[f]} |`);
}
console.log('');
console.log(`--weights='${JSON.stringify(fit)}'`);