/* Interactive companion — "Same Data, Different Losses, Same Circuits?"
   All charts client-side via Plotly from the paper's released metrics. */
(async function () {
  "use strict";

  const DATA = {};
  const files = ["meta", "cosine", "per_layer", "geometry", "accuracy", "seed_lr"];
  try {
    const loaded = await Promise.all(
      files.map((f) => fetch(`./data/${f}.json`).then((r) => {
        if (!r.ok) throw new Error(`${f}.json ${r.status}`);
        return r.json();
      }))
    );
    files.forEach((f, i) => (DATA[f] = loaded[i]));
  } catch (e) {
    document.querySelector("main").insertAdjacentHTML(
      "afterbegin",
      `<p style="color:#c0392b">Could not load data: ${e.message}. Serve this folder over HTTP (e.g. <code>python3 -m http.server</code>).</p>`
    );
    return;
  }

  const meta = DATA.meta;
  const ORDER = ["sft", "rft", "dft", "rift", "grpo", "dpo", "online_grpo", "online_dapo"];
  const SHORT = { sft: "SFT", rft: "RFT", dft: "DFT", rift: "RIFT", grpo: "GRPO", dpo: "DPO", online_grpo: "oGRPO", online_dapo: "oDAPO" };
  const C = meta.methodColors;

  // ---- flat print theme (ink on paper, no gradients) ----
  const css = (v) => getComputedStyle(document.body).getPropertyValue(v).trim();
  const INK = css("--ink") || "#211d18";
  const SOFT = css("--ink-soft") || "#50473c";
  const FAINT = css("--ink-faint") || "#897e6c";
  const ACCENT = css("--accent") || "#b23a23";
  const GRID = "rgba(31,27,22,0.10)";
  const AXIS = "rgba(31,27,22,0.55)";
  const MONO = "IBM Plex Mono, monospace";

  const CONFIG = { displayModeBar: false, responsive: true };
  function layout(extra) {
    return Object.assign({
      paper_bgcolor: "rgba(0,0,0,0)",
      plot_bgcolor: "rgba(0,0,0,0)",
      font: { family: MONO, color: SOFT, size: 11.5 },
      margin: { l: 58, r: 22, t: 14, b: 46 },
      hoverlabel: { bgcolor: INK, bordercolor: INK, font: { family: MONO, size: 11.5, color: "#f3efe4" } },
      xaxis: { gridcolor: GRID, zerolinecolor: AXIS, linecolor: AXIS, ticks: "outside", tickcolor: AXIS, ticklen: 4 },
      yaxis: { gridcolor: GRID, zerolinecolor: AXIS, linecolor: AXIS, ticks: "outside", tickcolor: AXIS, ticklen: 4 },
      legend: { orientation: "h", y: -0.2, font: { family: MONO, size: 11 } },
      showlegend: false,
    }, extra || {});
  }

  // ---- pair helpers ----
  const idIndex = (id) => ORDER.indexOf(id);
  const pairKey = (a, b) => [a, b].sort((x, y) => idIndex(x) - idIndex(y)).join("__");
  const pairLabel = (pk) => pk.split("__").map((m) => SHORT[m]).join(" · ");
  function labelToId(label) {
    const s = label.toLowerCase().replace(/\s+/g, "_");
    const map = { offline_grpo: "grpo", online_grpo: "online_grpo", online_dapo: "online_dapo" };
    return map[s] || s;
  }
  // stable per-pair palette — muted, risograph-ish, no neon
  const PAIR_PALETTE = ["#2f5d7c", "#bf4d2e", "#3f7a6d", "#c08a1e", "#7c5a86", "#5a6b3b",
    "#9c5a2e", "#456a8a", "#8a6d2e", "#7a3b3b", "#3f6e6e", "#874a6a",
    "#6a6a3b", "#9c3526", "#444038"];
  const masterPairs = Object.keys(DATA.per_layer.cosine);
  const pairColor = (pk) => PAIR_PALETTE[Math.max(0, masterPairs.indexOf(pk)) % PAIR_PALETTE.length];

  // =================================================================== header
  document.getElementById("venue").textContent = meta.venue;
  document.getElementById("title").textContent = meta.title;
  document.getElementById("subtitle").textContent = meta.subtitle;
  document.getElementById("abstract-text").textContent = meta.abstract;
  document.getElementById("footer-title").textContent = meta.title;
  for (const id of ["repolink", "repolink2"]) document.getElementById(id).href = meta.repo;
  if (meta.paper) document.getElementById("paperlink").href = meta.paper;
  document.querySelectorAll("figcaption[data-cap]").forEach((el) => {
    el.textContent = meta.captions[el.dataset.cap] || "";
  });

  // =================================================================== methods table
  (function methodsTable() {
    const t = document.getElementById("methods-table");
    t.innerHTML =
      `<thead><tr><th>Method</th><th>Loss</th><th>Negatives</th><th>Reward</th><th>Ref. policy</th></tr></thead><tbody>` +
      meta.methods.map((m) =>
        `<tr><td><span class="swatch" style="background:${C[m.id]}"></span>${m.name}</td>` +
        `<td>${m.loss.includes("×") || m.loss.includes("sg(") ? `<code>${m.loss}</code>` : m.loss}</td>` +
        `<td>${m.neg}</td><td>${m.reward}</td><td>${m.ref}</td></tr>`
      ).join("") + `</tbody>`;

    // objectives written out as math (KaTeX)
    const renderTeX = (tex) =>
      window.katex ? katex.renderToString(tex, { throwOnError: false, displayMode: false }) : tex;
    document.getElementById("objectives").innerHTML = meta.methods.map((m) =>
      `<div><dt><span class="swatch" style="background:${C[m.id]}"></span>${m.name}</dt>` +
      `<dd>${renderTeX(m.formula)}</dd></div>`
    ).join("");
    const ceEq = document.getElementById("ce-eq");
    if (ceEq && window.katex)
      ceEq.innerHTML = renderTeX("-\\,\\mathbb{E}\\,\\log \\pi_\\theta(y\\mid x)=\\operatorname{CE}(y,\\pi_\\theta)");
  })();

  // =================================================================== chips util
  function buildChips(host, pairs, selected, onChange) {
    host.classList.add("legend-toggles");
    host.innerHTML = "";
    pairs.forEach((pk) => {
      const b = document.createElement("button");
      b.className = "chip" + (selected.has(pk) ? " on" : "");
      b.style.setProperty("--swatch", pairColor(pk));
      b.innerHTML = `<span class="dot"></span>${pairLabel(pk)}`;
      b.addEventListener("click", () => {
        if (selected.has(pk)) {
          if (selected.size === 1) return; // keep at least one
          selected.delete(pk);
          b.classList.remove("on");
        } else { selected.add(pk); b.classList.add("on"); }
        onChange();
      });
      host.appendChild(b);
    });
  }

  // =================================================================== 1. cosine heatmap
  const cosineState = { size: 6 };
  // discrete, flat bands keyed to cosine thresholds (no continuous ramp).
  // normalized over zmin=-0.2..zmax=1; light warm steps so ink numerals stay legible.
  const BAND = ["#f0e9d7", "#e7d7bd", "#ddc29a", "#d3ab74", "#c8924f", "#bf7a39"];
  const STOP = [0, 0.25, 0.4167, 0.5833, 0.75, 0.9167, 1];
  const DISCRETE = [];
  BAND.forEach((c, i) => { DISCRETE.push([STOP[i], c]); DISCRETE.push([STOP[i + 1], c]); });
  function drawCosine() {
    const d = DATA.cosine;
    const labels = cosineState.size === 8 ? d.labels8 : d.labels6;
    const z = cosineState.size === 8 ? d.matrix8 : d.matrix6;
    Plotly.react("chart-cosine", [{
      type: "heatmap", z, x: labels, y: labels,
      zmin: -0.2, zmax: 1, colorscale: DISCRETE, showscale: false,
      xgap: 2, ygap: 2,
      text: z.map((row) => row.map((v) => v.toFixed(2))),
      texttemplate: "%{text}", textfont: { size: 12, family: MONO, color: INK },
      hovertemplate: "%{y} · %{x}  cos = %{z:.3f}<extra></extra>",
    }], layout({
      margin: { l: 112, r: 24, t: 8, b: 96 },
      xaxis: { tickangle: -40, gridcolor: "rgba(0,0,0,0)", zerolinecolor: "rgba(0,0,0,0)", linecolor: "rgba(0,0,0,0)", ticks: "", side: "bottom", automargin: true, tickfont: { family: MONO, size: 11, color: INK } },
      yaxis: { autorange: "reversed", gridcolor: "rgba(0,0,0,0)", zerolinecolor: "rgba(0,0,0,0)", linecolor: "rgba(0,0,0,0)", ticks: "", automargin: true, tickfont: { family: MONO, size: 11, color: INK } },
    }), CONFIG);
  }
  document.querySelectorAll("#cosine .seg-btn").forEach((b) =>
    b.addEventListener("click", () => {
      document.querySelectorAll("#cosine .seg-btn").forEach((x) => x.classList.remove("active"));
      b.classList.add("active");
      cosineState.size = +b.dataset.mx;
      drawCosine();
    }));
  drawCosine();
  // click a cell -> select that pair in the per-layer explorer
  document.getElementById("chart-cosine").on("plotly_click", (ev) => {
    const p = ev.points[0];
    const a = labelToId(p.y), b = labelToId(p.x);
    if (a === b) return;
    const pk = pairKey(a, b);
    if (!DATA.per_layer.cosine[pk]) return; // online pairs not in per-layer set
    perlayerSel.clear(); perlayerSel.add(pk);
    buildChips(document.getElementById("perlayer-pairs"), masterPairs, perlayerSel, drawPerlayer);
    drawPerlayer();
    document.getElementById("perlayer").scrollIntoView({ behavior: "smooth" });
  });

  // =================================================================== 2+3. per-layer cosine + shared layer slider
  const layers = DATA.per_layer.layers;
  const sharedLayer = { idx: 30 };
  const perlayerSel = new Set(["sft__rft", "sft__grpo", "sft__dpo"]);

  function lineTraces(source, pairs) {
    return [...pairs].map((pk) => ({
      type: "scatter", mode: "lines", name: pairLabel(pk),
      x: layers, y: source[pk],
      line: { color: pairColor(pk), width: 2, shape: "linear" },
      hovertemplate: `${pairLabel(pk)}  L%{x} = %{y:.3f}<extra></extra>`,
    }));
  }
  function vLine() {
    return [{
      type: "line", x0: sharedLayer.idx, x1: sharedLayer.idx, yref: "paper", y0: 0, y1: 1,
      line: { color: SOFT, width: 1, dash: "dot" },
    }];
  }
  function drawPerlayer() {
    Plotly.react("chart-perlayer", lineTraces(DATA.per_layer.cosine, perlayerSel),
      layout({
        showlegend: true,
        xaxis: { title: { text: "transformer block", font: { size: 12 } }, gridcolor: GRID, dtick: 5 },
        yaxis: { title: { text: "cosine of ΔW", font: { size: 12 } }, gridcolor: GRID, range: [-0.25, 1.05], zeroline: true },
        shapes: vLine(),
      }), CONFIG);
    updateReadout();
  }
  function updateReadout() {
    const L = layers[sharedLayer.idx];
    const parts = [...perlayerSel].map((pk) => `${pairLabel(pk)} = ${DATA.per_layer.cosine[pk][sharedLayer.idx].toFixed(3)}`);
    document.getElementById("layer-readout").textContent = `block ${L} · ` + parts.join("  ·  ");
  }
  buildChips(document.getElementById("perlayer-pairs"), masterPairs, perlayerSel, drawPerlayer);
  drawPerlayer();

  const slider = document.getElementById("layer-slider");
  slider.value = sharedLayer.idx;
  slider.addEventListener("input", () => {
    sharedLayer.idx = +slider.value;
    Plotly.relayout("chart-perlayer", { shapes: vLine() });
    Plotly.relayout("chart-cka", { shapes: vLine() });
    updateReadout();
  });

  // =================================================================== 4. CKA per-layer
  const ckaPairs = Object.keys(DATA.per_layer.cka);
  const ckaSel = new Set(["grpo__dpo", "sft__grpo", "sft__rift", "dft__dpo"].filter((p) => DATA.per_layer.cka[p]));
  if (ckaSel.size === 0) ckaPairs.slice(0, 3).forEach((p) => ckaSel.add(p));
  function drawCka() {
    Plotly.react("chart-cka", lineTraces(DATA.per_layer.cka, ckaSel),
      layout({
        showlegend: true,
        xaxis: { title: { text: "transformer block", font: { size: 12 } }, gridcolor: GRID, dtick: 5 },
        yaxis: { title: { text: "linear CKA", font: { size: 12 } }, gridcolor: GRID, range: [0.3, 1.03] },
        shapes: vLine(),
      }), CONFIG);
  }
  buildChips(document.getElementById("cka-pairs"), ckaPairs, ckaSel, drawCka);
  drawCka();

  // =================================================================== 5. output subspace (top-1 SVD u)
  const svdPairs = Object.keys(DATA.per_layer.svd_u);
  const svdSel = new Set(["sft__rft", "sft__dpo", "sft__grpo"]);
  function drawSvd() {
    Plotly.react("chart-svd", lineTraces(DATA.per_layer.svd_u, svdSel),
      layout({
        showlegend: true,
        xaxis: { title: { text: "transformer block", font: { size: 12 } }, gridcolor: GRID, dtick: 5 },
        yaxis: { title: { text: "top-1 output direction cosine", font: { size: 12 } }, gridcolor: GRID, range: [-0.1, 1.05], zeroline: true },
      }), CONFIG);
  }
  buildChips(document.getElementById("svd-pairs"), svdPairs, svdSel, drawSvd);
  drawSvd();

  // =================================================================== 6. principal angles
  (function drawAngles() {
    const pa = DATA.geometry.principal_angles;
    const keys = Object.keys(pa); // sft__X
    const others = keys.map((k) => k.split("__").find((m) => m !== "sft"));
    const labels = others.map((m) => SHORT[m]);
    Plotly.react("chart-angles", [
      { type: "bar", name: "top-1 angle", x: labels, y: keys.map((k) => pa[k].median_top1),
        marker: { color: others.map((m) => C[m]) },
        hovertemplate: "SFT · %{x}<br>top-1 = %{y:.1f}°<extra></extra>" },
      { type: "bar", name: "top-10 worst", x: labels, y: keys.map((k) => pa[k].median_top10_max),
        marker: { color: others.map((m) => C[m]), opacity: 0.32 },
        hovertemplate: "SFT · %{x}<br>top-10 worst = %{y:.1f}°<extra></extra>" },
    ], layout({
      barmode: "group", bargap: 0.35, showlegend: true,
      yaxis: { title: { text: "principal angle (degrees)", font: { size: 12 } }, gridcolor: GRID, range: [0, 95] },
      xaxis: { gridcolor: "rgba(0,0,0,0)" },
    }), CONFIG);
  })();

  // =================================================================== 7. update geometry (norm + rank)
  (function drawGeometry() {
    const g = DATA.geometry;
    const ms = ["sft", "rft", "dft", "rift", "grpo", "dpo"];
    const labels = ms.map((m) => SHORT[m]);
    Plotly.react("chart-geometry", [
      { type: "bar", name: "‖ΔW‖ (Frobenius)", x: labels, y: ms.map((m) => g.frobenius[m]),
        marker: { color: ms.map((m) => C[m]) },
        hovertemplate: "%{x}<br>‖ΔW‖ = %{y:.3f}<extra></extra>" },
      { type: "scatter", name: "effective rank", x: labels, y: ms.map((m) => g.effective_rank[m]),
        yaxis: "y2", mode: "lines+markers",
        line: { color: SOFT, width: 1.5, dash: "dot" }, marker: { size: 9, color: SOFT, symbol: "diamond" },
        hovertemplate: "%{x}<br>eff. rank = %{y:.1f}<extra></extra>" },
    ], layout({
      showlegend: true, bargap: 0.45,
      yaxis: { title: { text: "‖ΔW‖", font: { size: 12 } }, gridcolor: GRID, rangemode: "tozero" },
      yaxis2: { title: { text: "effective rank", font: { size: 12 } }, overlaying: "y", side: "right",
        gridcolor: "rgba(0,0,0,0)", rangemode: "tozero", showgrid: false },
      xaxis: { gridcolor: "rgba(0,0,0,0)" },
      legend: { orientation: "h", y: -0.16, font: { size: 12 } },
    }), CONFIG);
  })();

  // =================================================================== 8. mode connectivity
  (function drawLmc() {
    const g = DATA.geometry;
    const alphas = g.alphas;
    const traces = Object.entries(g.lmc).map(([k, v], i) => {
      const [a, b] = k.split("__");
      return {
        type: "scatter", mode: "lines+markers",
        name: `${SHORT[a]} → ${SHORT[b]}`, x: alphas, y: v,
        line: { color: PAIR_PALETTE[i % PAIR_PALETTE.length], width: 2, shape: "linear" },
        marker: { size: 5, symbol: "square" },
        hovertemplate: `${SHORT[a]}→${SHORT[b]}<br>α=%{x} · loss %{y:.2f}<extra></extra>`,
      };
    });
    Plotly.react("chart-lmc", traces, layout({
      showlegend: true,
      xaxis: { title: { text: "interpolation α  (0 = first adapter, 1 = second)", font: { size: 12 } }, gridcolor: GRID, dtick: 0.25 },
      yaxis: { title: { text: "masked-answer loss", font: { size: 12 } }, gridcolor: GRID, rangemode: "tozero" },
    }), CONFIG);
  })();

  // =================================================================== 9. accuracy
  const accState = { bench: "gsm8k" };
  function drawAccuracy() {
    const rows = DATA.accuracy.methods.filter((r) => r[accState.bench] != null);
    Plotly.react("chart-accuracy", [{
      type: "bar", orientation: "h",
      y: rows.map((r) => r.label).reverse(),
      x: rows.map((r) => r[accState.bench]).reverse(),
      marker: { color: rows.map((r) => C[r.method] || "#94A3B8").reverse() },
      text: rows.map((r) => `${r[accState.bench].toFixed(1)}%`).reverse(),
      textposition: "outside", textfont: { family: MONO, size: 12, color: INK },
      hovertemplate: "%{y}<br>%{x:.1f}%<extra></extra>",
    }], layout({
      margin: { l: 108, r: 40, t: 10, b: 42 },
      xaxis: { title: { text: "greedy pass@1 (%)", font: { size: 12 } }, gridcolor: GRID,
        range: [0, accState.bench === "gsm8k" ? 104 : 26] },
      yaxis: { gridcolor: "rgba(0,0,0,0)", automargin: true },
    }), CONFIG);
  }
  document.querySelectorAll("#accuracy .seg-btn").forEach((b) =>
    b.addEventListener("click", () => {
      document.querySelectorAll("#accuracy .seg-btn").forEach((x) => x.classList.remove("active"));
      b.classList.add("active");
      accState.bench = b.dataset.bench;
      drawAccuracy();
    }));
  drawAccuracy();

  // =================================================================== 10. seed / LR
  const seedState = { view: "cos" };
  function drawSeedLr() {
    const s = DATA.seed_lr;
    const lrs = s.lrs;
    if (seedState.view === "cos") {
      // SFT: raw cosine vs top-1 output direction, across LR (two seeds)
      const sft = s.seed_cosine.sft;
      const x = lrs.filter((lr) => sft[lr]);
      Plotly.react("chart-seedlr", [
        { type: "bar", name: "raw weight cosine", x, y: x.map((lr) => sft[lr].cos),
          marker: { color: C.sft }, hovertemplate: "lr %{x}<br>raw cosine = %{y:.3f}<extra></extra>" },
        { type: "scatter", name: "top-1 output dir (u₁)", x, y: x.map((lr) => sft[lr].top1_u),
          mode: "lines+markers", line: { color: C.dft, width: 2.4 }, marker: { size: 9 },
          hovertemplate: "lr %{x}<br>u₁ cosine = %{y:.3f}<extra></extra>" },
        { type: "scatter", name: "top-1 input dir (v₁)", x, y: x.map((lr) => sft[lr].top1_v),
          mode: "lines+markers", line: { color: C.online_dapo, width: 2, dash: "dot" }, marker: { size: 7 },
          hovertemplate: "lr %{x}<br>v₁ cosine = %{y:.3f}<extra></extra>" },
      ], layout({
        showlegend: true,
        yaxis: { title: { text: "seed-42 vs seed-123 cosine (SFT)", font: { size: 12 } }, gridcolor: GRID, range: [0, 1.05] },
        xaxis: { title: { text: "learning rate", font: { size: 12 } }, type: "category", gridcolor: "rgba(0,0,0,0)" },
      }), CONFIG);
    } else {
      // LR rotates dW: adjacent-LR cosine vs norm ratio (SFT vs DAPO)
      const dapo = s.dapo_lr_sensitivity;
      const steps = Object.keys(dapo);
      const labels = steps.map((k) => k.replace(/_to_/g, " → "));
      Plotly.react("chart-seedlr", [
        { type: "bar", name: "cosine between adjacent LRs (DAPO)", x: labels, y: steps.map((k) => dapo[k].cos),
          marker: { color: C.online_dapo }, hovertemplate: "%{x}<br>cosine = %{y:.3f}<extra></extra>" },
        { type: "scatter", name: "norm ratio (×)", x: labels, y: steps.map((k) => dapo[k].norm_ratio),
          yaxis: "y2", mode: "lines+markers", line: { color: SOFT, width: 1.5, dash: "dot" }, marker: { size: 9, symbol: "diamond" },
          hovertemplate: "%{x}<br>norm ×%{y:.1f}<extra></extra>" },
      ], layout({
        showlegend: true, bargap: 0.5,
        yaxis: { title: { text: "direction cosine", font: { size: 12 } }, gridcolor: GRID, range: [0, 1.05] },
        yaxis2: { title: { text: "norm ratio (×)", font: { size: 12 } }, overlaying: "y", side: "right", showgrid: false, rangemode: "tozero" },
        xaxis: { type: "category", gridcolor: "rgba(0,0,0,0)" },
      }), CONFIG);
    }
  }
  document.querySelectorAll("#seedlr .seg-btn").forEach((b) =>
    b.addEventListener("click", () => {
      document.querySelectorAll("#seedlr .seg-btn").forEach((x) => x.classList.remove("active"));
      b.classList.add("active");
      seedState.view = b.dataset.seed;
      drawSeedLr();
    }));
  drawSeedLr();

  // =================================================================== TOC scrollspy
  (function scrollspy() {
    const links = [...document.querySelectorAll("#toc a")];
    const map = new Map(links.map((a) => [a.getAttribute("href").slice(1), a]));
    const obs = new IntersectionObserver((entries) => {
      entries.forEach((e) => {
        if (e.isIntersecting) {
          links.forEach((l) => l.classList.remove("active"));
          const a = map.get(e.target.id);
          if (a) a.classList.add("active");
        }
      });
    }, { rootMargin: "-30% 0px -60% 0px" });
    document.querySelectorAll("main section").forEach((s) => obs.observe(s));
  })();
})();