// Lab: offline clustering experiments on window embeddings from lab-embed.mjs.
// Pipeline under test (mirrors the planned diar.js v2):
//   L2 → subtract recording mean → L2
//   stage A: average-linkage AHC at a FINE threshold (deliberate over-clustering)
//   fold: clusters with < minTotal s of turn time → nearest strong centroid
//   stage B: continue average-linkage merging of the strong clusters until
//            thrFinal (or until k clusters remain)
// Usage: node cluster-lab.mjs <emb.json> [--expect N] [--fine 0.35] [--final a,b,c]
//                             [--k N] [--minTotal 15] [--center 0|1] [--quiet]
import fs from 'node:fs';

const args = process.argv.slice(2);
const file = args[0];
const opt = (k, d) => { const i = args.indexOf('--' + k); return i >= 0 ? args[i + 1] : d; };
const has = k => args.includes('--' + k);
const EXPECT = +(opt('expect', 0));
const CENTER = opt('center', '1') === '1';
const FINE = +(opt('fine', 0.35));
const MIN_TOTAL = +(opt('minTotal', 15));
const KFIX = +(opt('k', 0));

const data = JSON.parse(fs.readFileSync(file, 'utf8'));
const turns = data.turns;
const wins = data.wins.filter(w => w.v);
const f = x => (x === undefined || x === null || Number.isNaN(x) ? 'NaN' : x.toFixed(3));

const norm = v => { let s = 0; for (const x of v) s += x * x; s = Math.sqrt(s) + 1e-9; return v.map(x => x / s); };
let X = wins.map(w => norm(w.v));
if (CENTER) {
  const dim = X[0].length, mu = new Array(dim).fill(0);
  for (const v of X) for (let i = 0; i < dim; i++) mu[i] += v[i] / X.length;
  X = X.map(v => norm(v.map((x, i) => x - mu[i])));
}
const dot = (a, b) => { let s = 0; for (let i = 0; i < a.length; i++) s += a[i] * b[i]; return s; };
const n = X.length;
const S = new Float32Array(n * n);
for (let i = 0; i < n; i++) { S[i * n + i] = 1; for (let j = i + 1; j < n; j++) { const s = dot(X[i], X[j]); S[i * n + j] = s; S[j * n + i] = s; } }

console.log(`# ${data.model.split('/').pop()} dim=${data.dim} windows=${n}/${data.wins.length} turns=${turns.length} center=${CENTER}`);

// ---- separability diagnostics ---------------------------------------------
const byTurn = new Map();
wins.forEach((w, i) => { (byTurn.get(w.ti) || byTurn.set(w.ti, []).get(w.ti)).push(i); });
const pos = [];
for (const idxs of byTurn.values())
  for (let a = 0; a < idxs.length; a++) for (let b = a + 1; b < idxs.length; b++)
    if (Math.abs(wins[idxs[a]].s - wins[idxs[b]].s) >= 3) pos.push(S[idxs[a] * n + idxs[b]]);
const rnd = [];
for (let k = 0; k < 40000; k++) {
  const i = (k * 7919) % n, j = (k * 104729 + 13) % n;
  if (i !== j && wins[i].ti !== wins[j].ti) rnd.push(S[i * n + j]);
}
const stat = a => {
  const m = a.reduce((x, y) => x + y, 0) / a.length;
  const sd = Math.sqrt(a.reduce((x, y) => x + (y - m) ** 2, 0) / a.length);
  const s = [...a].sort((x, y) => x - y);
  return { n: a.length, mean: m, sd, p05: s[(s.length * 0.05) | 0], p50: s[(s.length * 0.5) | 0], p95: s[(s.length * 0.95) | 0] };
};
const sp = stat(pos), sr = stat(rnd);
console.log(`# same-turn  mean=${f(sp.mean)} sd=${f(sp.sd)} p05=${f(sp.p05)}   cross-turn mean=${f(sr.mean)} sd=${f(sr.sd)} p95=${f(sr.p95)}   d'=${f((sp.mean - sr.mean) / Math.sqrt((sp.sd ** 2 + sr.sd ** 2) / 2))}`);

// ---- average-linkage AHC, optionally continuing from given groups -----------
function ahc(threshold, kStop = 0, initGroups = null) {
  const groups = initGroups ? initGroups.map(g => [...g]) : X.map((_, i) => [i]);
  const m = groups.length;
  const A = new Float32Array(m * m);
  for (let i = 0; i < m; i++) {
    A[i * m + i] = 1;
    for (let j = i + 1; j < m; j++) {
      let s = 0;
      for (const a of groups[i]) for (const b of groups[j]) s += S[a * n + b];
      s /= groups[i].length * groups[j].length;
      A[i * m + j] = s; A[j * m + i] = s;
    }
  }
  const cnt = new Int32Array(m); groups.forEach((g, i) => cnt[i] = g.length);
  const alive = new Uint8Array(m).fill(1);
  const bestJ = new Int32Array(m).fill(-1), bestS = new Float32Array(m).fill(-Infinity);
  const rescan = i => {
    let bj = -1, bs = -Infinity;
    for (let j = 0; j < m; j++) { if (j === i || !alive[j]) continue; const s = A[i * m + j]; if (s > bs) { bs = s; bj = j; } }
    bestJ[i] = bj; bestS[i] = bs;
  };
  for (let i = 0; i < m; i++) rescan(i);
  let k = m;
  while (k > 1) {
    let a = -1, bs = -Infinity;
    for (let i = 0; i < m; i++) if (alive[i] && bestS[i] > bs) { bs = bestS[i]; a = i; }
    if (a < 0) break;
    const b = bestJ[a]; if (b < 0) break;
    if (kStop > 0 ? k <= kStop : bs < threshold) break;
    const ca = cnt[a], cb = cnt[b];
    for (let c = 0; c < m; c++) {
      if (!alive[c] || c === a || c === b) continue;
      const v = (A[a * m + c] * ca + A[b * m + c] * cb) / (ca + cb);
      A[a * m + c] = v; A[c * m + a] = v;
    }
    cnt[a] = ca + cb; alive[b] = 0; k--;
    groups[a] = groups[a].concat(groups[b]); groups[b] = [];
    rescan(a);
    for (let i = 0; i < m; i++) if (alive[i] && (bestJ[i] === b || bestJ[i] === a)) rescan(i);
  }
  return groups.filter(g => g.length);
}

// ---- turn labels + per-cluster talk time -----------------------------------
function turnLabels(groups) {
  const of = new Int32Array(n).fill(-1);
  groups.forEach((g, gi) => g.forEach(i => of[i] = gi));
  const acc = new Map();
  wins.forEach((w, i) => {
    const mm = acc.get(w.ti) || acc.set(w.ti, new Map()).get(w.ti);
    mm.set(of[i], (mm.get(of[i]) || 0) + (w.e - w.s));
  });
  const lab = new Array(turns.length).fill(-1);
  for (const [ti, mm] of acc) { let bl = -1, bd = -1; for (const [l, d] of mm) if (d > bd) { bd = d; bl = l; } lab[ti] = bl; }
  const talk = new Map();
  turns.forEach((t, ti) => { if (lab[ti] >= 0) talk.set(lab[ti], (talk.get(lab[ti]) || 0) + (t.end - t.start)); });
  return { lab, talk };
}
const centroid = g => norm(g.reduce((acc, i) => (X[i].forEach((x, d) => acc[d] += x / g.length), acc), new Array(X[0].length).fill(0)));

// ---- the full pipeline -----------------------------------------------------
function run(tag, { fine = FINE, final = null, k = 0, minTotal = MIN_TOTAL }) {
  let groups = ahc(fine);
  const { talk } = turnLabels(groups);
  // fold weak clusters into the nearest strong centroid
  const strongIdx = groups.map((g, i) => i).filter(i => (talk.get(i) || 0) >= minTotal);
  let folded = 0;
  if (strongIdx.length && strongIdx.length < groups.length) {
    const cents = new Map(strongIdx.map(i => [i, centroid(groups[i])]));
    const merged = new Map(strongIdx.map(i => [i, [...groups[i]]]));
    groups.forEach((g, i) => {
      if (merged.has(i)) return;
      const c = centroid(g);
      let best = strongIdx[0], bs = -Infinity;
      for (const s of strongIdx) { const v = dot(c, cents.get(s)); if (v > bs) { bs = v; best = s; } }
      merged.get(best).push(...g); folded++;
    });
    groups = [...merged.values()];
  }
  const afterFold = groups.length;
  if (final !== null || k > 0) groups = ahc(final === null ? 0 : final, k, groups);
  const { lab, talk: t2 } = turnLabels(groups);
  const order = [...t2.entries()].sort((a, b) => b[1] - a[1]);
  const switches = lab.filter((l, i) => i > 0 && l >= 0 && lab[i - 1] >= 0 && l !== lab[i - 1]).length;
  const first = new Map(), last = new Map();
  turns.forEach((t, ti) => { const l = lab[ti]; if (l < 0) return; if (!first.has(l)) first.set(l, t.start); last.set(l, t.end); });
  console.log(`${tag}: fine=${f(fine)}${final !== null ? ` final=${f(final)}` : ''}${k ? ` k=${k}` : ''} → afterFine=${afterFold + folded} folded=${folded} speakers=${groups.length} switches=${switches}${EXPECT ? ` expect=${EXPECT}` : ''}`);
  if (!has('quiet')) console.log('   ' + order.map(([l, d]) => `${d.toFixed(0)}s[${(first.get(l) / 60).toFixed(1)}-${(last.get(l) / 60).toFixed(1)}]`).join(' '));
  return groups.length;
}

if (KFIX) run('pinned', { k: KFIX });
for (const t of (opt('final', '') ? opt('final').split(',').map(Number) : [0.10, 0.15, 0.20, 0.25, 0.30]))
  run('auto', { final: t });