/**
 * Project the corpus manifold into two dimensions.
 *
 * The 34 pieces of work live as points in a 384-dimensional embedding space.
 * That space has real structure — the chapters cluster, the theory pieces sit
 * between clusters, the papers sit at their own pole — and that structure is a
 * measured fact about what Joseph has written, not a decoration.
 *
 * This step flattens it to a plane while preserving the distances as faithfully
 * as two dimensions allow, so the shape can be drawn.
 *
 * Method: classical scaling for the initial configuration, then SMACOF stress
 * majorisation to refine it. SMACOF (Scaling by MAjorizing a COmplicated
 * Function, de Leeuw 1977) is guaranteed to decrease stress at every iteration,
 * so it cannot oscillate and needs no learning rate. For 34 points it converges
 * in well under a second.
 *
 * Distance is ANGULAR, not raw cosine: theta = arccos(sim) / pi. Angular
 * distance obeys the triangle inequality and cosine distance does not, and MDS
 * assumes a metric. Using the non-metric version produces a plausible-looking
 * layout with distorted neighbourhoods, which is the worst kind of wrong.
 *
 * Usage: node build/project.mjs <corpus> <embeddings> <out>
 */

import { readFile, writeFile } from 'node:fs/promises';

const SCALE = 32767;

/** Deterministic PRNG — the layout must be identical on every build. */
function mulberry32(seed) {
  let a = seed >>> 0;
  return () => {
    a = (a + 0x6d2b79f5) >>> 0;
    let t = a;
    t = Math.imul(t ^ (t >>> 15), 1 | t);
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}

function angularDistance(a, b) {
  let dot = 0; let na = 0; let nb = 0;
  for (let i = 0; i < a.length; i += 1) { dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; }
  const sim = Math.max(-1, Math.min(1, dot / (Math.sqrt(na) * Math.sqrt(nb))));
  return Math.acos(sim) / Math.PI;
}

/** Raw stress: the disagreement between wanted and achieved distances. */
function stress(D, X) {
  const n = X.length;
  let s = 0;
  for (let i = 0; i < n; i += 1) {
    for (let j = i + 1; j < n; j += 1) {
      const d = Math.hypot(X[i][0] - X[j][0], X[i][1] - X[j][1]);
      s += (d - D[i][j]) ** 2;
    }
  }
  return s;
}

/**
 * One SMACOF iteration: the Guttman transform.
 *   X' = (1/n) B(X) X
 * where B has off-diagonal -d_ij/dist_ij and rows summing to zero.
 */
function guttman(D, X) {
  const n = X.length;
  const out = Array.from({ length: n }, () => [0, 0]);
  for (let i = 0; i < n; i += 1) {
    let bx = 0; let by = 0; let bii = 0;
    for (let j = 0; j < n; j += 1) {
      if (i === j) continue;
      const dx = X[i][0] - X[j][0];
      const dy = X[i][1] - X[j][1];
      const dist = Math.hypot(dx, dy);
      // Guard the singularity: coincident points contribute nothing rather
      // than producing an infinity that poisons the whole configuration.
      const b = dist > 1e-9 ? -D[i][j] / dist : 0;
      bx += b * X[j][0];
      by += b * X[j][1];
      bii -= b;
    }
    out[i][0] = (bii * X[i][0] + bx) / n;
    out[i][1] = (bii * X[i][1] + by) / n;
  }
  return out;
}

async function main() {
  const [corpusFile, embFile, outFile] = process.argv.slice(2);

  const corpus = JSON.parse(await readFile(corpusFile, 'utf8'));
  const buf = await readFile(embFile);
  const view = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
  const count = view.getInt32(0, true);
  const dim = view.getInt32(4, true);
  if (count !== corpus.items.length) throw new Error(`embedding/corpus mismatch: ${count} vs ${corpus.items.length}`);

  const vecs = [];
  for (let i = 0; i < count; i += 1) {
    const v = new Float32Array(dim);
    for (let d = 0; d < dim; d += 1) v[d] = view.getInt16(8 + (i * dim + d) * 2, true) / SCALE;
    vecs.push(v);
  }

  // Target distances.
  const D = Array.from({ length: count }, () => new Array(count).fill(0));
  for (let i = 0; i < count; i += 1) {
    for (let j = i + 1; j < count; j += 1) {
      const d = angularDistance(vecs[i], vecs[j]);
      D[i][j] = d; D[j][i] = d;
    }
  }

  // Initial configuration: seeded random on a disc. Classical scaling would be
  // marginally better but needs an eigensolver; SMACOF from a random start
  // reaches the same basin for a set this size, and the seed keeps it stable.
  const rng = mulberry32(0x5eed);
  let X = Array.from({ length: count }, () => {
    const a = rng() * Math.PI * 2;
    const r = Math.sqrt(rng()) * 0.5;
    return [Math.cos(a) * r, Math.sin(a) * r];
  });

  let prev = stress(D, X);
  const s0 = prev;
  let iterations = 0;
  for (let it = 0; it < 3000; it += 1) {
    X = guttman(D, X);
    const s = stress(D, X);
    iterations = it + 1;
    // SMACOF decreases monotonically, so a tiny relative improvement means we
    // have converged rather than that we are stuck.
    if (prev - s < 1e-12) break;
    prev = s;
  }

  // Normalise into [0,1]^2 with a margin, so the renderer needs no magic numbers.
  const xs = X.map((p) => p[0]); const ys = X.map((p) => p[1]);
  const minX = Math.min(...xs); const maxX = Math.max(...xs);
  const minY = Math.min(...ys); const maxY = Math.max(...ys);
  const span = Math.max(maxX - minX, maxY - minY) || 1;
  const pad = 0.08;
  const points = X.map(([x, y], i) => ({
    id: corpus.items[i].id,
    topic: corpus.items[i].topic,
    title: corpus.items[i].title,
    // Centre the shorter axis so the manifold is not stretched to fill a square.
    x: +(pad + ((x - minX) - (maxX - minX - span) / 2) / span * (1 - 2 * pad)).toFixed(5),
    y: +(pad + ((y - minY) - (maxY - minY - span) / 2) / span * (1 - 2 * pad)).toFixed(5),
    // Mass drives how much the point contributes to the density field. Log of
    // word count: a 2800-word paper should weigh more than a 200-word note, but
    // not fourteen times more.
    mass: +Math.log1p(corpus.items[i].wordCount / 100).toFixed(4),
  }));

  const manifold = { count, iterations, stressStart: +s0.toFixed(4), stressEnd: +prev.toFixed(4), points };
  await writeFile(outFile, `${JSON.stringify(manifold)}\n`);

  // Report faithfulness: how well 2D preserves the 384-D neighbourhoods.
  let kept = 0; let total = 0;
  const K = 4;
  for (let i = 0; i < count; i += 1) {
    const hi = D[i].map((d, j) => ({ j, d })).filter((o) => o.j !== i).sort((a, b) => a.d - b.d).slice(0, K).map((o) => o.j);
    const lo = points.map((p, j) => ({ j, d: Math.hypot(points[i].x - p.x, points[i].y - p.y) }))
      .filter((o) => o.j !== i).sort((a, b) => a.d - b.d).slice(0, K).map((o) => o.j);
    kept += hi.filter((j) => lo.includes(j)).length;
    total += K;
  }

  console.log(`manifold: ${count} points, ${iterations} SMACOF iterations`);
  console.log(`stress ${s0.toFixed(3)} -> ${prev.toFixed(3)} (${((1 - prev / s0) * 100).toFixed(1)}% reduction)`);
  console.log(`top-${K} neighbourhood preserved: ${(kept / total * 100).toFixed(1)}%`);
  const byTopic = {};
  for (const p of points) (byTopic[p.topic] ??= []).push(p);
  for (const [t, ps] of Object.entries(byTopic)) {
    const cx = ps.reduce((s, p) => s + p.x, 0) / ps.length;
    const cy = ps.reduce((s, p) => s + p.y, 0) / ps.length;
    const spread = Math.sqrt(ps.reduce((s, p) => s + (p.x - cx) ** 2 + (p.y - cy) ** 2, 0) / ps.length);
    console.log(`  ${t.padEnd(10)} n=${String(ps.length).padStart(2)}  centroid (${cx.toFixed(2)}, ${cy.toFixed(2)})  spread ${spread.toFixed(3)}`);
  }
}

main().catch((e) => { console.error(e); process.exit(1); });