// Stage 04 — feed-forward MLP.
// Visualized as: vector expands up to 4×, passes through nonlinearity, projects back down.
function StageMLP({ tokens, activeIdx, progress, active, nLayers = 32, onExplain }) {
  const activeId = tokens[activeIdx]?.id ?? 0;
  const p = active ? progress : 0;

  // three columns of bars: input (d), expanded (4d), output (d)
  const nIn = 24, nHidden = 48, nOut = 24;

  const inputVals = React.useMemo(() => {
    const r = seededRand(activeId * 41 + 11);
    return Array.from({ length: nIn }, () => (r() - 0.5) * 2);
  }, [activeId]);

  const hiddenVals = React.useMemo(() => {
    const r = seededRand(activeId * 41 + 29);
    return Array.from({ length: nHidden }, () => Math.max(0, (r() - 0.3) * 2.4)); // GELU-ish, non-negative-ish
  }, [activeId]);

  const outputVals = React.useMemo(() => {
    const r = seededRand(activeId * 41 + 53);
    return Array.from({ length: nOut }, () => (r() - 0.5) * 2);
  }, [activeId]);

  // phases
  const phaseUp = clamp(p / 0.4, 0, 1);
  const phaseAct = clamp((p - 0.4) / 0.2, 0, 1);
  const phaseDown = clamp((p - 0.6) / 0.4, 0, 1);

  const barH = 80;
  const Col = ({ vals, max = 1, reveal, color = 'var(--ink)', absolute = false }) => (
    <div style={{
      display: 'grid', gridTemplateColumns: `repeat(${vals.length}, minmax(0, 1fr))`,
      gap: 1, height: barH, alignItems: 'end',
      border: '1px solid var(--line)', padding: '4px 6px', borderRadius: 4,
      background: 'rgba(0,0,0,0.25)',
    }}>
      {vals.map((v, i) => {
        const shown = reveal * vals.length > i;
        const h = absolute ? Math.min(1, v / max) : Math.min(1, Math.abs(v) / max);
        const sign = absolute ? 1 : Math.sign(v);
        return (
          <div key={i} style={{
            height: `${h * 100}%`,
            background: `color-mix(in oklch, ${sign < 0 ? 'var(--magenta)' : color} ${50 + h * 50}%, transparent)`,
            opacity: shown ? 1 : 0.08,
            boxShadow: shown ? `0 0 6px ${sign < 0 ? 'var(--magenta)' : color}` : 'none',
            transition: 'opacity 80ms',
          }} />
        );
      })}
    </div>
  );

  return (
    <Panel stage={4} title="Feed-forward × layers" subtitle={`residual · expand d→4d · GELU · project 4d→d · repeat ${nLayers}×`} active={active} progress={progress} onExplain={onExplain}>
      <div style={{ display: 'grid', gridTemplateColumns: 'minmax(0, 1fr) auto minmax(0, 1.6fr) auto minmax(0, 1fr)', gap: 10, alignItems: 'end' }}>
        <div>
          <div style={{ fontFamily: 'var(--mono)', fontSize: 9, color: 'var(--text-faint)', marginBottom: 4 }}>input · d</div>
          <Col vals={inputVals} reveal={1} />
        </div>
        <div style={{ fontFamily: 'var(--mono)', fontSize: 16, color: phaseUp > 0 ? 'var(--ink)' : 'var(--text-faint)', paddingBottom: 28, textShadow: phaseUp > 0 ? '0 0 8px var(--ink)' : 'none', transition: 'color 200ms' }}>W↑</div>
        <div>
          <div style={{ fontFamily: 'var(--mono)', fontSize: 9, color: 'var(--text-faint)', marginBottom: 4, display: 'flex', justifyContent: 'space-between' }}>
            <span>hidden · 4d (GELU)</span>
            <span style={{ color: phaseAct > 0 ? 'var(--amber)' : 'var(--text-faint)', textShadow: phaseAct > 0 ? '0 0 6px var(--amber)' : 'none' }}>nonlinearity</span>
          </div>
          <Col vals={hiddenVals} max={1.5} reveal={phaseAct} color={phaseAct > 0.2 ? 'var(--amber)' : 'var(--ink)'} absolute />
        </div>
        <div style={{ fontFamily: 'var(--mono)', fontSize: 16, color: phaseDown > 0 ? 'var(--ink)' : 'var(--text-faint)', paddingBottom: 28, textShadow: phaseDown > 0 ? '0 0 8px var(--ink)' : 'none', transition: 'color 200ms' }}>W↓</div>
        <div>
          <div style={{ fontFamily: 'var(--mono)', fontSize: 9, color: 'var(--text-faint)', marginBottom: 4 }}>output · d (+ residual)</div>
          <Col vals={outputVals} reveal={phaseDown} />
        </div>
      </div>

      {/* layer stack indicator */}
      <div style={{ marginTop: 10, display: 'flex', alignItems: 'center', gap: 8, fontFamily: 'var(--mono)', fontSize: 9, color: 'var(--text-faint)' }}>
        <span>layer</span>
        <div style={{ display: 'flex', gap: 2, flex: 1 }}>
          {Array.from({ length: nLayers }).map((_, i) => {
            const reached = active && (progress * nLayers) > i;
            return (
              <div key={i} style={{
                flex: 1, height: 4, borderRadius: 1,
                background: reached ? 'var(--ink)' : 'rgba(140,255,160,0.12)',
                boxShadow: reached ? '0 0 4px var(--ink)' : 'none',
                transition: 'background 120ms, box-shadow 120ms',
              }} />
            );
          })}
        </div>
        <span>{String(Math.min(nLayers, Math.ceil((active ? progress : 0) * nLayers))).padStart(2, '0')} / {nLayers}</span>
      </div>
    </Panel>
  );
}
window.StageMLP = StageMLP;