// Stage 03 — multi-head self-attention.
// Visualized as: a grid of attention heatmaps (one per head) with a subtle
// 3D float + mouse-tracked parallax so "parallel heads" reads as literal
// stacked planes rather than flat swatches.
function StageAttention({ tokens, activeIdx, progress, active, nHeads = 8, onExplain }) {
  const T = tokens.length || 1;
  const activeId = tokens[activeIdx]?.id ?? 0;

  // Parallax: track mouse position inside the heads grid (-1..1 on each axis).
  // Throttled via rAF so we never re-render faster than the browser's frame.
  const gridRef = React.useRef(null);
  const [parallax, setParallax] = React.useState({ x: 0, y: 0 });
  React.useEffect(() => {
    const el = gridRef.current; if (!el) return;
    let raf = 0;
    let pending = { x: 0, y: 0 };
    const onMove = (e) => {
      const r = el.getBoundingClientRect();
      pending = {
        x: ((e.clientX - r.left) / r.width - 0.5) * 2,
        y: ((e.clientY - r.top)  / r.height - 0.5) * 2,
      };
      if (!raf) raf = requestAnimationFrame(() => { raf = 0; setParallax(pending); });
    };
    const onLeave = () => setParallax({ x: 0, y: 0 });
    el.addEventListener('mousemove', onMove);
    el.addEventListener('mouseleave', onLeave);
    return () => {
      el.removeEventListener('mousemove', onMove);
      el.removeEventListener('mouseleave', onLeave);
      if (raf) cancelAnimationFrame(raf);
    };
  }, []);

  // Compute a causal attention pattern per head (deterministic, stylized)
  const heads = React.useMemo(() => {
    const hs = [];
    for (let h = 0; h < nHeads; h++) {
      const r = seededRand(activeId * 31 + h * 97 + 3);
      const rows = [];
      for (let i = 0; i < T; i++) {
        const row = new Array(T).fill(0);
        let sum = 0;
        for (let j = 0; j <= i; j++) {
          // Different heads "prefer" different positions
          let bias;
          if (h % 4 === 0) bias = Math.exp(-0.6 * (i - j));                      // recent
          else if (h % 4 === 1) bias = 1;                                         // uniform
          else if (h % 4 === 2) bias = j === 0 ? 2.5 : 0.3;                       // BOS-heavy
          else bias = Math.exp(-0.25 * Math.abs(i - j - Math.floor(i / 2)));      // mid-range
          const v = bias * (0.6 + r() * 0.8);
          row[j] = v; sum += v;
        }
        for (let j = 0; j < T; j++) row[j] = sum > 0 ? row[j] / sum : 0;
        rows.push(row);
      }
      hs.push(rows);
    }
    return hs;
  }, [T, activeId, nHeads]);

  // reveal controls: first build QKV (0..0.33), then softmax (0.33..0.66), then weighted sum (0.66..1)
  const p = active ? progress : 0;
  const phaseQKV = clamp(p / 0.33, 0, 1);
  const phaseSoft = clamp((p - 0.33) / 0.33, 0, 1);
  const phaseSum = clamp((p - 0.66) / 0.34, 0, 1);

  const cellSize = Math.min(14, Math.max(6, Math.floor(150 / Math.max(T, 4))));

  return (
    <Panel stage={3} title="Multi-head self-attention" subtitle={`${nHeads} heads · causal mask · softmax(QKᵀ/√d)·V`} active={active} progress={progress} onExplain={onExplain}>
      <div style={{ display: 'grid', gridTemplateColumns: '160px minmax(0, 1fr)', gap: 14 }}>
        {/* Left column: Q/K/V strips */}
        <div style={{ display: 'flex', flexDirection: 'column', gap: 6 }}>
          {['Q', 'K', 'V'].map((label, li) => (
            <div key={label} style={{
              border: '1px solid var(--line)', borderRadius: 4, padding: '5px 7px',
              background: 'rgba(0,0,0,0.25)',
            }}>
              <div style={{ display: 'flex', justifyContent: 'space-between', fontFamily: 'var(--mono)', fontSize: 9, color: 'var(--text-faint)', marginBottom: 3 }}>
                <span>{label} · projection</span>
                <span>{label === 'Q' ? 'question' : label === 'K' ? 'key' : 'value'}</span>
              </div>
              <div style={{ display: 'flex', gap: 1, height: 10 }}>
                {Array.from({ length: 32 }).map((_, i) => {
                  const r = mulberry32(activeId * 13 + li * 17 + i)();
                  const appeared = phaseQKV * 32 > i;
                  return (
                    <div key={i} style={{
                      flex: 1, height: '100%',
                      background: `color-mix(in oklch, var(--ink) ${Math.round(r * 100)}%, transparent)`,
                      opacity: appeared ? 1 : 0.08,
                      transition: 'opacity 80ms',
                    }} />
                  );
                })}
              </div>
            </div>
          ))}
          <div style={{ fontFamily: 'var(--mono)', fontSize: 9, color: 'var(--text-faint)', marginTop: 2 }}>
            split across {nHeads} heads →
          </div>
        </div>

        {/* Right: heatmap grid of heads, now on a 3D stage. Each card
            floats with a slight translateZ keyed to its position; mouse
            parallax rotates the whole group so the stack reads as three-
            dimensional without losing legibility of individual heads. */}
        <div
          ref={gridRef}
          data-attention-heads-grid
          style={{
            display: 'grid',
            gridTemplateColumns: `repeat(${nHeads / 2}, minmax(0, 1fr))`,
            gap: 10,
            perspective: '1200px',
            perspectiveOrigin: '50% 50%',
            isolation: 'isolate',
          }}>
          {heads.map((head, hi) => {
            const revealed = phaseSoft > hi / nHeads;
            // Per-card depth offset: heads at the corners sit slightly
            // further back so the grid looks like a fan rather than a plane.
            const col = hi % (nHeads / 2);
            const rowIdx = Math.floor(hi / (nHeads / 2));
            const colOff = (col - (nHeads / 2 - 1) / 2) * 6;  // lateral Z
            const rowOff = (rowIdx - 0.5) * 8;                 // row Z
            const zBase = -12 - Math.abs(colOff) - Math.abs(rowOff);
            // Parallax rotate (shared across grid). Magnitude bumped so
            // the depth reads clearly — still well under nausea threshold.
            const rx = -parallax.y * 6;
            const ry =  parallax.x * 8;
            return (
              <div key={hi} style={{
                border: `1px solid ${revealed ? 'var(--line-strong)' : 'var(--line)'}`,
                borderRadius: 3, padding: 4,
                background: 'rgba(0,0,0,0.25)',
                opacity: revealed ? 1 : 0.35,
                transform: `rotateX(${rx}deg) rotateY(${ry}deg) translateZ(${zBase}px)`,
                transformStyle: 'preserve-3d',
                backfaceVisibility: 'hidden',
                transition: 'opacity 200ms, border-color 200ms, transform 200ms cubic-bezier(0.2, 0.8, 0.2, 1)',
                boxShadow: `0 ${4 - rowIdx * 2}px ${10 - rowIdx * 4}px color-mix(in oklch, var(--ink) ${Math.round(6 + Math.abs(rowOff))}%, transparent)`,
              }}>
                <div style={{ display: 'flex', justifyContent: 'space-between', fontFamily: 'var(--mono)', fontSize: 8, color: 'var(--text-faint)', marginBottom: 3 }}>
                  <span>head {hi}</span>
                  <span>{hi % 4 === 0 ? 'recent' : hi % 4 === 1 ? 'uniform' : hi % 4 === 2 ? 'first' : 'mid'}</span>
                </div>
                <svg viewBox={`0 0 ${T * cellSize} ${T * cellSize}`} style={{ width: '100%', height: 'auto', display: 'block' }}>
                  {head.map((row, i) => row.map((v, j) => {
                    if (j > i) return null; // causal mask
                    const w = v;
                    const isActiveRow = i === activeIdx;
                    const color = isActiveRow ? 'var(--amber)' : 'var(--ink)';
                    return (
                      <rect key={`${i}-${j}`} x={j * cellSize} y={i * cellSize} width={cellSize - 0.5} height={cellSize - 0.5}
                        fill={color} opacity={Math.min(1, w * 1.8) * (isActiveRow ? 1 : 0.85)}
                        style={{ transition: 'opacity 200ms' }} />
                    );
                  }))}
                </svg>
              </div>
            );
          })}
        </div>
      </div>

      {/* bottom: flow lines from past tokens to current */}
      <div style={{
        marginTop: 10, position: 'relative', height: 28,
        borderTop: '1px dashed var(--line-strong)', paddingTop: 10,
      }}>
        <svg style={{ position: 'absolute', inset: 0, width: '100%', height: '100%' }} preserveAspectRatio="none" viewBox="0 0 100 28">
          {tokens.map((_, i) => {
            if (i > activeIdx) return null;
            const x1 = (i + 0.5) / T * 100;
            const x2 = (activeIdx + 0.5) / T * 100;
            // aggregate attention weight from heads at (activeIdx, i)
            const w = heads.reduce((acc, h) => acc + (h[activeIdx]?.[i] || 0), 0) / heads.length;
            return (
              <path key={i} d={`M${x1} 2 Q ${(x1 + x2) / 2} 22, ${x2} 26`}
                fill="none" stroke="var(--ink)" strokeWidth={0.3 + w * 1.4}
                opacity={phaseSum * Math.min(1, w * 3)}
                style={{ filter: `drop-shadow(0 0 ${w * 2}px var(--ink))` }} />
            );
          })}
        </svg>
      </div>
    </Panel>
  );
}
window.StageAttention = StageAttention;