// ── Tab Navigation ──
document.querySelectorAll('.tab').forEach(btn => {
  btn.addEventListener('click', () => {
    document.querySelectorAll('.tab').forEach(b => b.classList.remove('active'));
    document.querySelectorAll('.panel').forEach(p => p.classList.remove('active'));
    btn.classList.add('active');
    document.getElementById(btn.dataset.tab).classList.add('active');
  });
});
document.querySelectorAll('.next-btn').forEach(btn => {
  btn.addEventListener('click', () => {
    document.querySelector(`.tab[data-tab="${btn.dataset.next}"]`).click();
    window.scrollTo({ top: 0, behavior: 'smooth' });
  });
});

// ── Helpers ──
function softmax(scores) {
  const max = Math.max(...scores);
  const exps = scores.map(s => Math.exp(s - max));
  const sum = exps.reduce((a, b) => a + b, 0);
  return exps.map(e => e / sum);
}

// ══════════════════════════════════════════
// TAB 1: Softmax Temperature Demo
// ══════════════════════════════════════════
const rawScores = [1.2, 0.5, 3.8, 0.9, 1.5]; // index 2 is the "target"
const scoreLabels = ['slot 0', 'slot 1', 'slot 2', 'slot 3', 'slot 4'];

function renderSoftmaxBars(temp) {
  const scaled = rawScores.map(s => s * temp);
  const weights = softmax(scaled);
  const maxW = Math.max(...weights);
  const container = document.getElementById('softmaxBars');
  container.innerHTML = '';
  weights.forEach((w, i) => {
    const col = document.createElement('div');
    col.className = 'bar-col';
    const wrapper = document.createElement('div');
    wrapper.className = 'bar-wrapper';
    const bar = document.createElement('div');
    bar.className = 'bar' + (w === maxW ? ' winner' : '');
    bar.style.height = (w * 130) + 'px';
    wrapper.appendChild(bar);
    const val = document.createElement('div');
    val.className = 'bar-value';
    val.textContent = (w * 100).toFixed(1) + '%';
    const lbl = document.createElement('div');
    lbl.className = 'bar-label';
    lbl.textContent = scoreLabels[i] + (i === 2 ? ' ★' : '');
    col.append(val, wrapper, lbl);
    container.appendChild(col);
  });
  const insight = document.getElementById('tempInsight');
  if (temp < 5) insight.textContent = 'At low temperature, attention is spread out — fuzzy, not useful for exact computation.';
  else if (temp < 20) insight.textContent = 'Getting sharper! The target slot is winning, but there\'s still leakage to other slots.';
  else insight.textContent = 'Nearly 100% on the target. The softmax now acts like an exact array read — this is how weights produce deterministic lookups.';
}

document.getElementById('tempSlider').addEventListener('input', e => {
  const v = +e.target.value;
  document.getElementById('tempVal').textContent = v;
  renderSoftmaxBars(v);
});
renderSoftmaxBars(1);

// ══════════════════════════════════════════
// TAB 2: Memory Lookup via Attention
// ══════════════════════════════════════════
const memValues = [42, 17, 99, 8, 55, 73];
let queryTarget = 2;

function makeColVec(values, cls, label) {
  // Returns a DOM element showing a column vector with bracket notation
  const wrap = document.createElement('div');
  wrap.className = 'col-vec ' + cls;
  wrap.innerHTML = '<div class="bracket"></div><div class="bracket-r"></div>';
  values.forEach(v => {
    const cell = document.createElement('div');
    cell.className = 'cell';
    cell.textContent = v;
    wrap.appendChild(cell);
  });
  if (label) {
    const lbl = document.createElement('div');
    lbl.className = 'vec-label';
    lbl.innerHTML = label;
    wrap.appendChild(lbl);
  }
  return wrap;
}

function renderMemory() {
  // Memory slots
  const container = document.getElementById('memorySlots');
  container.innerHTML = '';
  memValues.forEach((v, i) => {
    const slot = document.createElement('div');
    slot.className = 'mem-slot' + (i === queryTarget ? ' active' : '');
    slot.innerHTML = `<div class="idx">addr ${i}</div><div class="val">${v}</div>`;
    slot.addEventListener('click', () => { queryTarget = i; renderMemory(); });
    container.appendChild(slot);
  });

  // Query vector
  const qEl = document.getElementById('queryVec');
  qEl.innerHTML = '';
  const qVec = makeColVec([queryTarget, 1], 'query', `q`);
  const qNote = document.createElement('span');
  qNote.style.cssText = 'font-size:0.82rem;color:var(--dim);margin-left:12px';
  qNote.innerHTML = `= (i, 1) where i = ${queryTarget} &nbsp;← "I want to read address ${queryTarget}"`;
  qEl.appendChild(qVec);
  qEl.appendChild(qNote);

  // Key vectors + dot products
  const vecEl = document.getElementById('vecColumns');
  vecEl.innerHTML = '';

  const scores = memValues.map((_, j) => 2 * queryTarget * j - j * j);
  const maxScore = Math.max(...scores);
  const minScore = Math.min(...scores);
  const scoreRange = maxScore - minScore || 1;
  const weights = softmax(scores.map(s => s * 10));

  memValues.forEach((val, j) => {
    const isWin = j === queryTarget;
    const k = [2 * j, -(j * j)];

    const group = document.createElement('div');
    group.className = 'vec-group' + (isWin ? ' winner' : '');

    // Column vector
    const vec = makeColVec(k, isWin ? 'winner' : '', `k<sub>${j}</sub>`);
    group.appendChild(vec);

    // Dot product computation
    const comp = document.createElement('div');
    comp.className = 'dot-computation';
    comp.innerHTML = `${queryTarget}×${k[0]} + 1×${k[1] >= 0 ? k[1] : '(' + k[1] + ')'}`;
    group.appendChild(comp);

    // Score + weight
    const dpLine = document.createElement('div');
    dpLine.className = 'dot-product-line';
    dpLine.innerHTML = `<span class="dp-val">${scores[j]}</span><span class="dp-weight">${(weights[j] * 100).toFixed(1)}%</span>`;
    group.appendChild(dpLine);

    // Mini bar
    const bar = document.createElement('div');
    bar.className = 'dp-bar-mini';
    bar.style.width = Math.max(2, ((scores[j] - minScore) / scoreRange) * 60) + 'px';
    group.appendChild(bar);

    // Value stored
    const valLabel = document.createElement('div');
    valLabel.style.cssText = `font-size:0.7rem;margin-top:4px;color:${isWin ? 'var(--gold)' : 'var(--dim)'};font-family:monospace`;
    valLabel.textContent = `val=${val}`;
    group.appendChild(valLabel);

    vecEl.appendChild(group);

    // Add "·" or "=" separator between groups (except last)
    if (j < memValues.length - 1) {
      const sep = document.createElement('div');
      sep.style.cssText = 'align-self:center;color:var(--border);font-size:1.2rem;padding:0 2px';
      sep.textContent = '';
      vecEl.appendChild(sep);
    }
  });

  // Read result
  document.getElementById('readResult').innerHTML =
    `<strong style="color:var(--gold)">Read result:</strong> mem[${queryTarget}] = <strong>${memValues[queryTarget]}</strong> &nbsp;— key k<sub>${queryTarget}</sub> gets score <strong>${maxScore}</strong> (softmax weight ${(weights[queryTarget] * 100).toFixed(2)}%), all others are penalized by −(i−j)²`;
}

renderMemory();

// ══════════════════════════════════════════
// TAB 2b: Side-by-Side Comparison
// ══════════════════════════════════════════
const sbsWords = ['The', 'cake', 'delicious', 'was', 'very'];
// Simulated high-dim embeddings (4D slice for display) — designed to show semantic similarity spread
const sbsTradKeys = [
  [0.2, -0.1,  0.8,  0.3],  // The
  [0.9,  0.7,  0.1, -0.2],  // cake
  [0.8,  0.9, -0.1,  0.3],  // delicious
  [0.1, -0.3,  0.7,  0.5],  // was
  [0.3,  0.1,  0.2,  0.9],  // very
];
// Queries are similar to the target but with overlap to neighbors (semantic similarity)
const sbsTradQueries = [
  [0.3, -0.2,  0.9,  0.2],  // attending to "The"
  [0.8,  0.6,  0.2, -0.1],  // attending to "cake"
  [0.7,  0.8,  0.0,  0.4],  // attending to "delicious"
  [0.2, -0.2,  0.8,  0.4],  // attending to "was"
  [0.4,  0.2,  0.1,  0.8],  // attending to "very"
];

let sbsTarget = 2;

function makeSbsColVec(values, cls, label) {
  const wrap = document.createElement('div');
  wrap.className = 'col-vec sbs-vec ' + cls;
  wrap.innerHTML = '<div class="bracket"></div><div class="bracket-r"></div>';
  values.forEach(v => {
    const cell = document.createElement('div');
    cell.className = 'cell';
    cell.textContent = typeof v === 'number' ? (Number.isInteger(v) ? v : v.toFixed(1)) : v;
    wrap.appendChild(cell);
  });
  if (label) {
    const lbl = document.createElement('div');
    lbl.className = 'vec-label';
    lbl.innerHTML = label;
    wrap.appendChild(lbl);
  }
  return wrap;
}

function renderSBS() {
  const target = sbsTarget;
  document.getElementById('sbsTargetLabel').textContent = sbsWords[target];

  // ── Traditional side ──
  const tradQEl = document.getElementById('sbsTradQ');
  tradQEl.innerHTML = '';
  const tradQLabel = document.createElement('span');
  tradQLabel.style.cssText = 'font-size:0.75rem;color:var(--dim);margin-right:4px';
  tradQLabel.textContent = 'query:';
  tradQEl.appendChild(tradQLabel);
  tradQEl.appendChild(makeSbsColVec(sbsTradQueries[target], 'query', 'q'));

  const tradKeysEl = document.getElementById('sbsTradKeys');
  tradKeysEl.innerHTML = '';

  const tradScores = sbsTradKeys.map(k =>
    k.reduce((sum, ki, d) => sum + ki * sbsTradQueries[target][d], 0)
  );
  const tradWeights = softmax(tradScores.map(s => s * 4)); // moderate temperature
  const tradMaxW = Math.max(...tradWeights);

  sbsWords.forEach((word, j) => {
    const grp = document.createElement('div');
    const isTop = tradWeights[j] === tradMaxW;
    grp.className = 'sbs-key-group' + (isTop ? ' trad-winner' : '');

    const wordEl = document.createElement('div');
    wordEl.className = 'sbs-word';
    wordEl.textContent = word;
    grp.appendChild(wordEl);

    grp.appendChild(makeSbsColVec(sbsTradKeys[j], '', `k<sub>${j}</sub>`));

    const score = document.createElement('div');
    score.className = 'sbs-score';
    score.textContent = tradScores[j].toFixed(2);
    grp.appendChild(score);

    const weight = document.createElement('div');
    weight.className = 'sbs-weight';
    weight.textContent = (tradWeights[j] * 100).toFixed(1) + '%';
    grp.appendChild(weight);

    const bar = document.createElement('div');
    bar.className = 'sbs-weight-bar';
    bar.style.width = (tradWeights[j] / tradMaxW * 50) + 'px';
    grp.appendChild(bar);

    tradKeysEl.appendChild(grp);
  });

  const tradResult = document.getElementById('sbsTradResult');
  const topTrad = tradWeights.map((w, i) => ({w, i})).sort((a, b) => b.w - a.w);
  tradResult.innerHTML = `<span style="color:var(--accent2)">Output:</span> blend of "${sbsWords[topTrad[0].i]}" (${(topTrad[0].w*100).toFixed(0)}%) + "${sbsWords[topTrad[1].i]}" (${(topTrad[1].w*100).toFixed(0)}%) + others<br><span style="color:var(--dim)">→ A fuzzy mix of semantically related tokens</span>`;

  // ── Lookup side ──
  const lookupQEl = document.getElementById('sbsLookupQ');
  lookupQEl.innerHTML = '';
  const lookupQLabel = document.createElement('span');
  lookupQLabel.style.cssText = 'font-size:0.75rem;color:var(--dim);margin-right:4px';
  lookupQLabel.textContent = 'query:';
  lookupQEl.appendChild(lookupQLabel);
  lookupQEl.appendChild(makeSbsColVec([target, 1], 'query', 'q'));

  const lookupKeysEl = document.getElementById('sbsLookupKeys');
  lookupKeysEl.innerHTML = '';

  const lookupScores = sbsWords.map((_, j) => 2 * target * j - j * j);
  const lookupWeights = softmax(lookupScores.map(s => s * 10));
  const lookupMaxW = Math.max(...lookupWeights);

  sbsWords.forEach((word, j) => {
    const grp = document.createElement('div');
    const isWin = lookupWeights[j] === lookupMaxW;
    grp.className = 'sbs-key-group' + (isWin ? ' winner' : '');

    const wordEl = document.createElement('div');
    wordEl.className = 'sbs-word';
    wordEl.textContent = `addr ${j}`;
    grp.appendChild(wordEl);

    grp.appendChild(makeSbsColVec([2 * j, -(j * j)], isWin ? 'winner' : '', `k<sub>${j}</sub>`));

    const score = document.createElement('div');
    score.className = 'sbs-score';
    score.textContent = lookupScores[j];
    grp.appendChild(score);

    const weight = document.createElement('div');
    weight.className = 'sbs-weight';
    weight.textContent = (lookupWeights[j] * 100).toFixed(1) + '%';
    grp.appendChild(weight);

    const bar = document.createElement('div');
    bar.className = 'sbs-weight-bar';
    bar.style.width = (lookupWeights[j] / (lookupMaxW || 1) * 50) + 'px';
    grp.appendChild(bar);

    lookupKeysEl.appendChild(grp);
  });

  const lookupResult = document.getElementById('sbsLookupResult');
  lookupResult.innerHTML = `<span style="color:var(--gold)">Output:</span> value at addr ${target} with <strong>${(lookupWeights[target] * 100).toFixed(2)}%</strong> weight<br><span style="color:var(--dim)">→ An exact read of one specific address</span>`;
}

document.getElementById('sbsTargetSlider').addEventListener('input', e => {
  sbsTarget = +e.target.value;
  renderSBS();
});
renderSBS();

// ══════════════════════════════════════════
// TAB 3: Parabola Visualization
// ══════════════════════════════════════════
function drawParabola(queryIdx) {
  const canvas = document.getElementById('parabolaCanvas');
  const ctx = canvas.getContext('2d');
  const W = canvas.width, H = canvas.height;
  ctx.clearRect(0, 0, W, H);

  const n = 8;
  const pad = 40;
  const xScale = (W - 2 * pad) / (2 * (n - 1));
  const maxJ2 = (n - 1) * (n - 1);
  const yScale = (H - 2 * pad) / maxJ2;

  // Axes
  ctx.strokeStyle = '#2a2a44';
  ctx.lineWidth = 1;
  ctx.beginPath();
  ctx.moveTo(pad, H - pad);
  ctx.lineTo(W - pad, H - pad);
  ctx.moveTo(pad, H - pad);
  ctx.lineTo(pad, pad);
  ctx.stroke();

  ctx.fillStyle = '#666680';
  ctx.font = '10px monospace';
  ctx.fillText('2j →', W - pad - 20, H - pad + 15);
  ctx.fillText('−j²', pad - 5, pad - 5);

  // Parabola curve
  ctx.strokeStyle = '#333355';
  ctx.lineWidth = 1.5;
  ctx.beginPath();
  for (let j = 0; j < n; j++) {
    const x = pad + (2 * j) * xScale;
    const y = H - pad - (j * j) * yScale;
    j === 0 ? ctx.moveTo(x, y) : ctx.lineTo(x, y);
  }
  ctx.stroke();

  // Points
  for (let j = 0; j < n; j++) {
    const x = pad + (2 * j) * xScale;
    const y = H - pad - (j * j) * yScale;
    ctx.beginPath();
    ctx.arc(x, y, j === queryIdx ? 8 : 5, 0, Math.PI * 2);
    ctx.fillStyle = j === queryIdx ? '#ffd54f' : '#4fc3f7';
    ctx.fill();
    ctx.fillStyle = '#999';
    ctx.font = '10px monospace';
    ctx.fillText(`j=${j}`, x - 8, y + 18);
  }

  // Query direction arrow
  const qx = pad + (2 * queryIdx) * xScale;
  const qy = H - pad - (queryIdx * queryIdx) * yScale;
  ctx.strokeStyle = '#ff7043';
  ctx.lineWidth = 2;
  ctx.setLineDash([4, 3]);
  ctx.beginPath();
  ctx.moveTo(pad + W / 4, H - pad);
  ctx.lineTo(qx, qy);
  ctx.stroke();
  ctx.setLineDash([]);

  ctx.fillStyle = '#ff7043';
  ctx.font = 'bold 11px sans-serif';
  ctx.fillText(`q=(${queryIdx},1)`, pad + W / 4 - 15, H - pad + 15);
}

function drawScores(queryIdx) {
  const canvas = document.getElementById('scoresCanvas');
  const ctx = canvas.getContext('2d');
  const W = canvas.width, H = canvas.height;
  ctx.clearRect(0, 0, W, H);

  const n = 8;
  const pad = 40;
  const barW = (W - 2 * pad) / n - 4;

  const scores = [];
  for (let j = 0; j < n; j++) {
    scores.push(2 * queryIdx * j - j * j);
  }
  const maxS = Math.max(...scores);
  const minS = Math.min(...scores);
  const range = maxS - minS || 1;

  // Axes
  ctx.strokeStyle = '#2a2a44';
  ctx.lineWidth = 1;
  ctx.beginPath();
  ctx.moveTo(pad, H - pad);
  ctx.lineTo(W - pad, H - pad);
  ctx.stroke();

  ctx.fillStyle = '#666680';
  ctx.font = '10px monospace';
  ctx.fillText('j →', W - pad - 15, H - pad + 15);
  ctx.fillText('score', pad - 5, pad + 10);

  // Bars
  for (let j = 0; j < n; j++) {
    const x = pad + j * ((W - 2 * pad) / n) + 2;
    const h = ((scores[j] - minS) / range) * (H - 2 * pad - 20);
    const y = H - pad - h;
    ctx.fillStyle = j === queryIdx ? '#ffd54f' : '#4fc3f7';
    ctx.fillRect(x, y, barW, h);
    ctx.fillStyle = '#999';
    ctx.font = '9px monospace';
    ctx.fillText(j.toString(), x + barW / 2 - 3, H - pad + 12);
    ctx.fillStyle = j === queryIdx ? '#ffd54f' : '#aaa';
    ctx.font = '9px monospace';
    ctx.fillText(scores[j].toString(), x + barW / 2 - 6, y - 4);
  }

  document.getElementById('parabolaInsight').textContent =
    `Index ${queryIdx} gets the highest score (${maxS}). The penalty −(i−j)² ensures only the exact match wins.`;
}

document.getElementById('queryIdxSlider').addEventListener('input', e => {
  const v = +e.target.value;
  document.getElementById('queryIdxVal').textContent = v;
  drawParabola(v);
  drawScores(v);
});

// Defer canvas drawing until tab is visible
const tab3Observer = new MutationObserver(() => {
  if (document.getElementById('tab3').classList.contains('active')) {
    drawParabola(3);
    drawScores(3);
    tab3Observer.disconnect();
  }
});
tab3Observer.observe(document.getElementById('tab3'), { attributes: true, attributeFilter: ['class'] });

// ══════════════════════════════════════════
// TAB 4: Write-Read Trace Demo
// ══════════════════════════════════════════
const wrSteps = [
  {
    type: 'write',
    instr: 'i32.const 3',
    desc: 'Push 3 onto stack',
    token: { label: 'token 0', instr: 'const 3', k: 'k=[2, −1]', v: 'v=3', addr: 1, val: 3 },
    action: '<div class="wr-step-title">WRITE: i32.const 3</div>The model emits a new trace token. W<sub>K</sub> maps it to key <strong style="color:var(--accent)">[2, −1]</strong> (stack depth 1 on the parabola). W<sub>V</sub> extracts value <strong style="color:var(--green)">3</strong>. The token now sits in the sequence — that\'s the write. No memory chip needed.',
  },
  {
    type: 'write',
    instr: 'i32.const 5',
    desc: 'Push 5 onto stack',
    token: { label: 'token 1', instr: 'const 5', k: 'k=[4, −4]', v: 'v=5', addr: 2, val: 5 },
    action: '<div class="wr-step-title">WRITE: i32.const 5</div>Another trace token emitted. Key <strong style="color:var(--accent)">[4, −4]</strong> (stack depth 2). Value <strong style="color:var(--green)">5</strong>. Now two tokens sit in the sequence = two stack entries.',
  },
  {
    type: 'read',
    instr: 'i32.add (read operands)',
    desc: 'Read top two stack values',
    readTargets: [1, 0],  // indices into tokens array
    action: '<div class="wr-step-title">READ: i32.add needs operands</div>The add instruction needs the top two stack values. The stack head produces query <strong style="color:var(--warn)">q=[2, 1]</strong> → attention scans all past tokens\' keys → finds token 1 (score = 2×2×2 − 4 = 4, highest) → retrieves value <strong style="color:var(--gold)">5</strong>. Then query <strong style="color:var(--warn)">q=[1, 1]</strong> → finds token 0 → retrieves <strong style="color:var(--gold)">3</strong>. No memory was accessed — just attention over past tokens.',
  },
  {
    type: 'write',
    instr: 'i32.add (result)',
    desc: 'Push result 8',
    token: { label: 'token 2', instr: 'add → 8', k: 'k=[2, −1]', v: 'v=8', addr: 1, val: 8 },
    shadowIdx: 0,  // token 0 gets overshadowed (same addr)
    action: '<div class="wr-step-title">WRITE: push result 8</div>The FFN computed 3 + 5 = 8. A new token is emitted with key <strong style="color:var(--accent)">[2, −1]</strong> (stack depth 1 — the stack shrank by 1). Value <strong style="color:var(--green)">8</strong>.<br><br>Notice: token 0 also had key [2, −1] (depth 1, value 3). But token 2 is <em>later</em> in the sequence, so the parabola trick gives it a higher score. <strong>The old value 3 is overshadowed, not erased.</strong>',
  },
  {
    type: 'read',
    instr: 'output',
    desc: 'Read top of stack',
    readTargets: [2],
    action: '<div class="wr-step-title">READ: output top of stack</div>Query <strong style="color:var(--warn)">q=[1, 1]</strong> for stack depth 1. Both token 0 and token 2 have key [2, −1], but token 2 is later → higher score → attention returns <strong style="color:var(--gold)">8</strong> (not the old 3). Output: <strong>8</strong>. ✓',
  },
];

let wrStep = 0;
let wrTokens = [];
let wrShadowed = new Set();

function wrReset() {
  wrStep = 0;
  wrTokens = [];
  wrShadowed = new Set();
  renderWR();
}

function wrStepExec() {
  if (wrStep >= wrSteps.length) return;
  const step = wrSteps[wrStep];
  if (step.type === 'write') {
    wrTokens.push({ ...step.token, isNew: true });
    if (step.shadowIdx !== undefined) wrShadowed.add(step.shadowIdx);
  }
  wrStep++;
  renderWR();
  // Clear "new" flag after animation
  setTimeout(() => {
    wrTokens.forEach(t => t.isNew = false);
    // Don't re-render, just let CSS animation finish
  }, 700);
}

function renderWR() {
  const traceEl = document.getElementById('wrTrace');
  traceEl.innerHTML = '';

  const step = wrStep > 0 ? wrSteps[wrStep - 1] : null;
  const isRead = step && step.type === 'read';
  const readSet = isRead ? new Set(step.readTargets) : new Set();

  if (wrTokens.length === 0) {
    traceEl.innerHTML = '<div style="color:var(--dim);font-size:0.82rem;padding:1rem;text-align:center">No tokens yet. The sequence is empty — no memory exists.<br>Click Step to emit the first trace token.</div>';
  }

  wrTokens.forEach((tok, i) => {
    const div = document.createElement('div');
    let cls = 'wr-token';
    if (tok.isNew) cls += ' new-token';
    if (isRead && readSet.has(i)) cls += ' found';
    if (wrShadowed.has(i)) cls += ' shadowed';
    div.className = cls;

    let inner = `<div class="wr-tok-label">${tok.label}</div>`;
    inner += `<div class="wr-tok-instr">${tok.instr}</div>`;
    inner += `<div class="wr-tok-kv"><span class="wr-k">${tok.k}</span><br><span class="wr-v">${tok.v}</span></div>`;
    if (isRead && readSet.has(i)) {
      inner += `<div class="wr-read-arrow">↑ READ</div>`;
    }
    if (wrShadowed.has(i) && !readSet.has(i)) {
      inner += `<div style="font-size:0.6rem;color:var(--warn);margin-top:2px">overshadowed</div>`;
    }
    div.innerHTML = inner;
    traceEl.appendChild(div);
  });

  const actionEl = document.getElementById('wrAction');
  if (step) {
    actionEl.innerHTML = step.action;
  } else {
    actionEl.innerHTML = '<span style="color:var(--dim)">Click Step to begin execution. Watch how each token becomes a memory cell.</span>';
  }
}

document.getElementById('wrStep').addEventListener('click', wrStepExec);
document.getElementById('wrReset').addEventListener('click', wrReset);
renderWR();

// ══════════════════════════════════════════
// TAB 4: Stack Machine Step-Through
// ══════════════════════════════════════════
const smProgram = [
  { op: 'i32.const', arg: 3, desc: 'Push 3' },
  { op: 'i32.const', arg: 5, desc: 'Push 5' },
  { op: 'i32.add', arg: null, desc: 'Pop two, push sum' },
  { op: 'i32.const', arg: 10, desc: 'Push 10' },
  { op: 'i32.sub', arg: null, desc: 'Pop two, subtract' },
  { op: 'output', arg: null, desc: 'Output top of stack' },
  { op: 'halt', arg: null, desc: 'Stop' },
];

let smState = { ip: 0, stack: [], output: null, done: false, headLog: [] };

function smReset() {
  smState = { ip: 0, stack: [], output: null, done: false, headLog: [] };
  renderSM();
}

function smStepExec() {
  if (smState.done) return;
  const instr = smProgram[smState.ip];
  const heads = [];

  // IP head: cumulative sum
  heads.push({ name: 'IP Head', action: `sum of deltas → IP = ${smState.ip}`, detail: `query: uniform avg × t = ${smState.ip}` });

  if (instr.op === 'i32.const') {
    smState.stack.push(instr.arg);
    heads.push({ name: 'Stack Head', action: `WRITE ${instr.arg} at depth ${smState.stack.length}`, detail: `key=(${2 * smState.stack.length}, -${smState.stack.length ** 2})  val=${instr.arg}` });
    heads.push({ name: 'FFN (ALU)', action: 'passthrough (no arithmetic)', detail: 'gate=1, val=input' });
  } else if (instr.op === 'i32.add') {
    const b = smState.stack.pop(), a = smState.stack.pop();
    const r = a + b;
    heads.push({ name: 'Stack Head ×2', action: `READ depth ${smState.stack.length + 2} → ${b}, depth ${smState.stack.length + 1} → ${a}`, detail: `q=(${smState.stack.length + 2},1) → ${b};  q=(${smState.stack.length + 1},1) → ${a}` });
    heads.push({ name: 'FFN (ALU)', action: `${a} + ${b} = ${r}`, detail: `ReLU gate selects ADD path` });
    smState.stack.push(r);
  } else if (instr.op === 'i32.sub') {
    const b = smState.stack.pop(), a = smState.stack.pop();
    const r = a - b;
    heads.push({ name: 'Stack Head ×2', action: `READ depth ${smState.stack.length + 2} → ${b}, depth ${smState.stack.length + 1} → ${a}`, detail: `q=(${smState.stack.length + 2},1) → ${b};  q=(${smState.stack.length + 1},1) → ${a}` });
    heads.push({ name: 'FFN (ALU)', action: `${a} - ${b} = ${r}`, detail: `ReLU gate selects SUB path` });
    smState.stack.push(r);
  } else if (instr.op === 'output') {
    const top = smState.stack[smState.stack.length - 1];
    smState.output = top;
    heads.push({ name: 'Stack Head', action: `READ top (depth ${smState.stack.length}) → ${top}`, detail: `q=(${smState.stack.length},1) → ${top}` });
  } else if (instr.op === 'halt') {
    smState.done = true;
    heads.push({ name: 'Control Head', action: 'HALT detected', detail: 'opcode matches halt pattern' });
  }

  smState.headLog = heads;
  smState.ip++;
  renderSM();
}

function renderSM() {
  // Program listing
  const progEl = document.getElementById('smProgram');
  progEl.innerHTML = '<h4>Program</h4>';
  smProgram.forEach((instr, i) => {
    const line = document.createElement('div');
    line.className = 'instr-line' + (i === smState.ip - 1 && !smState.done ? ' current' : '') + (i < smState.ip - 1 ? ' done' : '') + (i === smState.ip - 1 && smState.done ? ' current' : '');
    line.textContent = `${i}: ${instr.op}${instr.arg !== null ? ' ' + instr.arg : ''}`;
    progEl.appendChild(line);
  });

  // State
  const stateEl = document.getElementById('smState');
  stateEl.innerHTML = '<h4>VM State</h4>';
  const rows = [
    ['IP', smState.ip >= smProgram.length ? 'HALT' : smState.ip],
    ['Stack depth', smState.stack.length],
    ['Output', smState.output !== null ? smState.output : '—'],
  ];
  rows.forEach(([l, v]) => {
    const row = document.createElement('div');
    row.className = 'state-row';
    row.innerHTML = `<span class="label">${l}</span><span class="value">${v}</span>`;
    stateEl.appendChild(row);
  });

  const stackLabel = document.createElement('div');
  stackLabel.style.cssText = 'font-size:0.75rem;color:var(--dim);margin-top:8px;margin-bottom:4px';
  stackLabel.textContent = 'Stack (top ↑):';
  stateEl.appendChild(stackLabel);

  [...smState.stack].reverse().forEach((v, i) => {
    const item = document.createElement('div');
    item.className = 'stack-item' + (i === 0 ? ' top' : '');
    item.textContent = v;
    stateEl.appendChild(item);
  });

  // Heads
  const headsEl = document.getElementById('smHeads');
  headsEl.innerHTML = '<h4>Attention Heads Active</h4>';
  if (smState.headLog.length === 0) {
    headsEl.innerHTML += '<div style="color:var(--dim);font-size:0.8rem">Click Step to begin</div>';
  }
  smState.headLog.forEach(h => {
    const div = document.createElement('div');
    div.className = 'head-info';
    div.innerHTML = `<div class="head-name">${h.name}</div><div class="head-action">${h.action}</div><div class="head-detail">${h.detail}</div>`;
    headsEl.appendChild(div);
  });
}

document.getElementById('smStep').addEventListener('click', smStepExec);
document.getElementById('smReset').addEventListener('click', smReset);
renderSM();

// IP demo (mini)
(function () {
  const el = document.querySelector('#ipDemo .ip-trace');
  const deltas = [1, 1, 1, 1, -2, 1, 1];
  let sum = 0;
  deltas.forEach(d => {
    sum += d;
    const cell = document.createElement('div');
    cell.className = 'ip-cell';
    cell.innerHTML = `<div class="delta">${d > 0 ? '+' : ''}${d}</div><div class="sum">IP=${sum}</div>`;
    el.appendChild(cell);
  });
})();

// Stack demo (mini)
(function () {
  const el = document.querySelector('#stackDemo .stack-vis');
  [3, 5, 8].forEach(v => {
    const item = document.createElement('div');
    item.className = 'sv-item';
    item.textContent = v;
    el.appendChild(item);
  });
})();

// ══════════════════════════════════════════
// TAB 5: Full Execution Trace
// ══════════════════════════════════════════
const feProgram = [
  { op: 'i32.const', bytes: '03 00 00 00', desc: 'Push 3 onto stack' },
  { op: 'i32.const', bytes: '05 00 00 00', desc: 'Push 5 onto stack' },
  { op: 'i32.add', bytes: '00 00 00 00', desc: 'Pop 3 and 5, push 8' },
  { op: 'output', bytes: '00 00 00 00', desc: 'Output top of stack' },
];

const feTraceSteps = [
  { tok: '03 00 00 00', meta: 'commit(+1,sts=1,bt=0)', detail: 'IP Head reads instruction 0 → i32.const. Stack Head writes 3 at depth 1. Stack: [3]' },
  { tok: '05 00 00 00', meta: 'commit(+1,sts=2,bt=0)', detail: 'IP Head reads instruction 1 → i32.const. Stack Head writes 5 at depth 2. Stack: [3, 5]' },
  { tok: '08 00 00 00', meta: 'commit(-1,sts=1,bt=0)', detail: 'IP Head reads instruction 2 → i32.add. Stack Head reads depth 2 → 5, depth 1 → 3. FFN computes 3+5=8. Writes 8 at depth 1. Stack: [8]' },
  { tok: 'out(08)', meta: '', detail: 'IP Head reads instruction 3 → output. Stack Head reads top → 8. Output token emitted.' },
  { tok: 'halt', meta: '', detail: 'Program complete. All computation happened inside the transformer\'s forward pass.' },
];

let feStep = 0;

function feReset() {
  feStep = 0;
  renderFE();
}

function feStepExec() {
  if (feStep < feTraceSteps.length) feStep++;
  renderFE();
}

function feRunAll() {
  feStep = feTraceSteps.length;
  renderFE();
}

function renderFE() {
  // Program
  const progEl = document.getElementById('feProgram');
  progEl.innerHTML = '<h4>WASM Program</h4>';
  feProgram.forEach((instr, i) => {
    const line = document.createElement('div');
    line.className = 'instr-line' + (i < feStep ? ' done' : '') + (i === feStep - 1 ? ' current' : '');
    line.textContent = `${instr.op} ${instr.bytes}`;
    progEl.appendChild(line);
  });

  // Trace
  const traceEl = document.getElementById('feTrace');
  traceEl.innerHTML = '<h4>Execution Trace (tokens)</h4>';
  for (let i = 0; i < feStep; i++) {
    const t = feTraceSteps[i];
    const div = document.createElement('div');
    div.className = 'trace-token' + (i === feStep - 1 ? ' new' : '');
    div.innerHTML = `<span class="tok">${t.tok}</span><span class="meta">${t.meta}</span>`;
    traceEl.appendChild(div);
  }
  if (feStep === 0) {
    traceEl.innerHTML += '<div style="color:var(--dim);font-size:0.8rem;padding:4px">Click Step to generate trace tokens...</div>';
  }

  // Detail
  const detailEl = document.getElementById('feDetail');
  detailEl.innerHTML = '<h4>What Happened (weight level)</h4>';
  if (feStep > 0) {
    const d = feTraceSteps[feStep - 1];
    const div = document.createElement('div');
    div.style.cssText = 'font-size:0.82rem;line-height:1.6;color:var(--text)';
    div.textContent = d.detail;
    detailEl.appendChild(div);

    // Visual: which heads fired
    const heads = [];
    if (feStep <= 3) heads.push({ name: 'IP Head', color: 'var(--accent)' });
    if (feStep <= 4) heads.push({ name: 'Stack Head', color: 'var(--green)' });
    if (feStep === 3) heads.push({ name: 'FFN (ALU)', color: 'var(--gold)' });
    if (feStep === 5) heads.push({ name: 'Control', color: 'var(--warn)' });

    const hDiv = document.createElement('div');
    hDiv.style.cssText = 'margin-top:12px;display:flex;gap:6px;flex-wrap:wrap';
    heads.forEach(h => {
      const chip = document.createElement('span');
      chip.style.cssText = `font-size:0.75rem;padding:3px 10px;border-radius:12px;border:1px solid ${h.color};color:${h.color}`;
      chip.textContent = h.name;
      hDiv.appendChild(chip);
    });
    detailEl.appendChild(hDiv);
  } else {
    detailEl.innerHTML += '<div style="color:var(--dim);font-size:0.8rem">Each step shows which attention heads fire and what they compute.</div>';
  }
}

document.getElementById('feStep').addEventListener('click', feStepExec);
document.getElementById('feReset').addEventListener('click', feReset);
document.getElementById('feRunAll').addEventListener('click', feRunAll);
renderFE();