frontend/src/components/LogitLens.js

import React, { useState, useEffect, useMemo, useCallback } from 'react';
import Plot from '../utils/PlotlyWrapper';
import { Send, Loader2, Info } from 'lucide-react';

const BACKEND_URL = process.env.REACT_APP_BACKEND_URL || '';

export const LogitLens = () => {
  const [prompt, setPrompt] = useState('When Mary and John went to the store, John gave a bottle of milk to');
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState(null);
  const [lensData, setLensData] = useState(null);

  const [modelStatus, setModelStatus] = useState(() => {
    return sessionStorage.getItem('cs_model_status') || 'demo';
  });

  useEffect(() => {
    const handleStatusChange = (e) => {
      setModelStatus(e.detail);
    };
    window.addEventListener('cs_model_status_change', handleStatusChange);
    return () => {
      window.removeEventListener('cs_model_status_change', handleStatusChange);
    };
  }, []);

  const handleFetchLens = useCallback(async (customPrompt) => {
    const targetPrompt = customPrompt || prompt;
    if (!targetPrompt.trim()) return;

    setLoading(true);
    setError(null);
    try {
      const res = await fetch(`${BACKEND_URL}/api/logit-lens`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ prompt: targetPrompt.trim() }),
      });
      if (!res.ok) {
        const errData = await res.json().catch(() => ({}));
        throw new Error(errData.detail || `Request failed (${res.status})`);
      }
      const data = await res.json();
      setLensData(data);
    } catch (e) {
      setError(e.message || 'Failed to project residual stream through Logit Lens');
    } finally {
      setLoading(false);
    }
  }, [prompt]);

  // Initial load
  useEffect(() => {
    handleFetchLens();
  }, [handleFetchLens]);

  const heatmapTrace = useMemo(() => {
    if (!lensData) return null;
    return {
      z: lensData.target_token_probs,
      x: lensData.tokens,
      y: lensData.layers.map(l => `Layer ${l}`),
      type: 'heatmap',
      colorscale: [
        [0, '#0C0F1A'],
        [0.2, '#1E2B45'],
        [0.5, '#6b21a8'], // Deep purple
        [0.8, '#9B59F5'], // Glowing purple
        [1.0, '#00D9C0'], // Neon Teal peak
      ],
      hoverongaps: false,
      hovertemplate: '<b>Token:</b> %{x}<br><b>%{y}</b><br><b>Target Prob:</b> %{z:.4f}<extra></extra>',
      showscale: true,
      colorbar: {
        title: 'Target Prob',
        titlefont: { color: '#8A9BC4', size: 10, family: 'JetBrains Mono' },
        tickfont: { color: '#8A9BC4', size: 9, family: 'JetBrains Mono' },
      }
    };
  }, [lensData]);

  const annotations = useMemo(() => {
    if (!lensData) return [];
    const ann = [];
    for (let l = 0; l < lensData.layers.length; l++) {
      for (let p = 0; p < lensData.tokens.length; p++) {
        const top1 = lensData.top_predictions[l][p][0]; // [token_str, prob]
        if (top1) {
          ann.push({
            x: lensData.tokens[p],
            y: `Layer ${l}`,
            text: `${top1[0]}<br><span style="font-size:9px;opacity:0.6">${parseFloat(top1[1]).toFixed(2)}</span>`,
            showarrow: false,
            font: {
              family: 'JetBrains Mono, monospace',
              size: 10,
              color: parseFloat(top1[1]) > 0.4 ? '#E8EEF8' : '#8A9BC4',
            }
          });
        }
      }
    }
    return ann;
  }, [lensData]);

  const chartLayout = useMemo(() => {
    if (!lensData) return {};
    return {
      paper_bgcolor: 'transparent',
      plot_bgcolor: 'transparent',
      margin: { l: 80, r: 20, t: 30, b: 60 },
      xaxis: {
        tickfont: { color: '#8A9BC4', size: 11, family: 'Space Grotesk' },
        gridcolor: '#1E2B4533',
        zeroline: false,
      },
      yaxis: {
        tickfont: { color: '#8A9BC4', size: 11, family: 'Space Grotesk' },
        gridcolor: '#1E2B4533',
        zeroline: false,
      },
      annotations: annotations,
    };
  }, [lensData, annotations]);

  return (
    <section id="logit-lens" className="scroll-mt-section" data-testid="section-logit-lens" style={{ padding: '80px 0', borderTop: '1px solid #1E2B4533' }}>
      <div className="section-container">
        <div className="mb-2">
          <span className="badge-violet">Phase 2: Logit Lens</span>
          <span style={{ marginLeft: 8 }} className="badge-green">
            GPT-2 Small Active
          </span>
        </div>
        
        <h2 style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 600, fontSize: 30, color: '#E8EEF8', marginBottom: 8 }}>
          Residual Stream Logit Lens
        </h2>
        <p style={{ fontSize: 15, color: '#8A9BC4', maxWidth: 680, marginBottom: 32, lineHeight: 1.75 }}>
          Project intermediate hidden states directly into vocabulary space layer-by-layer using the unembed matrix $W_U$. Watch the model's intermediate "best guesses" shift from early syntax representations to late semantic target resolution.
        </p>

        {/* Input Interface */}
        <div className="research-card mb-8">
          <div className="flex items-center gap-3">
            <input
              type="text"
              value={prompt}
              onChange={(e) => setPrompt(e.target.value)}
              disabled={loading}
              placeholder="Input custom IOI sequence..."
              style={{
                flex: 1,
                background: '#080B12',
                border: '1px solid #1E2B45',
                borderRadius: 8,
                padding: '12px 16px',
                color: '#E8EEF8',
                fontSize: 14,
                fontFamily: "'Space Grotesk', sans-serif",
                outline: 'none',
              }}
            />
            <button
              onClick={() => handleFetchLens()}
              disabled={loading}
              className="btn-primary flex items-center gap-2"
              style={{ padding: '12px 24px' }}
            >
              {loading ? <Loader2 className="animate-spin" size={16} /> : <Send size={16} />}
              Project Lens
            </button>
          </div>
          {error && (
            <div style={{ marginTop: 12, display: 'flex', alignItems: 'center', gap: 8, color: '#FF5063', fontSize: 13 }}>
              <span style={{ fontWeight: 600 }}>Error:</span> {error}
            </div>
          )}
        </div>

        {/* Dynamic heatmap */}
        {lensData && (
          <div className="research-card mb-8">
            <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', marginBottom: 16 }}>
              <h3 style={{ fontFamily: "'Space Grotesk', sans-serif", fontSize: 16, color: '#E8EEF8', fontWeight: 600 }}>
                Layer Vocabulary Projection Map
              </h3>
              <div style={{ fontSize: 12, color: '#8A9BC4', display: 'flex', alignItems: 'center', gap: 6 }}>
                <Info size={14} /> Heatmap color scale shows target name token probability (<strong>{lensData.io_name}</strong>)
              </div>
            </div>
            
            <div className="plotly-container" style={{ overflowX: 'auto' }}>
              <div style={{ minWidth: 700 }}>
                <Plot
                  data={[heatmapTrace]}
                  layout={chartLayout}
                  config={{ displayModeBar: false, responsive: true }}
                />
              </div>
            </div>
          </div>
        )}

        {/* Research Interpretability Insights */}
        <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
          <div className="research-card-accent" style={{ borderLeftColor: '#9B59F5' }}>
            <h4 style={{ fontSize: 14, color: '#E8EEF8', fontWeight: 600, marginBottom: 8 }}>
              Mechanistic Interpretation
            </h4>
            <p style={{ fontSize: 13, color: '#8A9BC4', lineHeight: 1.6 }}>
              Projecting hidden states via the logit lens illustrates that the transformer does not solve sequence completion instantly. In early layers (0–4), representations are highly syntactic (predicting punctuation or connector words). The correct name prediction (e.g. <strong>{lensData?.io_name}</strong>) emerges sharply around layers 7–9, directly coinciding with S-inhibition and name-mover attention interventions.
            </p>
          </div>
          <div className="research-card-accent" style={{ borderLeftColor: '#00D9C0' }}>
            <h4 style={{ fontSize: 14, color: '#E8EEF8', fontWeight: 600, marginBottom: 8 }}>
              Technical Specifications
            </h4>
            <p style={{ fontSize: 13, color: '#8A9BC4', lineHeight: 1.6 }}>
              For layer {"$L$"}, the logit distribution is extracted by projecting the residual stream representation through the final layer normalization followed by the standard unembed weight: {"$\\text{logits}_L = \\text{LN}(x_L) \\cdot W_U$"}. Projecting onto intermediate layers provides a highly precise approximation of the model's internal coordinate updates.
            </p>
          </div>
        </div>
      </div>
    </section>
  );
};