import React, { useState, useMemo, useCallback, useEffect } from 'react';
import { Zap, RotateCcw, Info, Loader2, Play } from 'lucide-react';
import circuitData from '../data/circuit_heads.json';

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

export const CircuitExplorer = () => {
  const [selectedGroup, setSelectedGroup] = useState(null);
  const [hoveredHead, setHoveredHead] = useState(null);
  
  // Ablation Knockout States
  const [knockedOutHeads, setKnockedOutHeads] = useState([]); // List of "L.H" strings, e.g. "5.1", "6.9"
  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 [knockoutResult, setKnockoutResult] = useState(null);
  const [ablationMode, setAblationMode] = useState('zero'); // 'zero' or 'mean'

  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 allHeads = useMemo(() => {
    return circuitData.groups.flatMap(g => g.heads.map(h => ({ ...h, groupId: g.id, groupName: g.name, groupColor: g.color, groupDesc: g.description })));
  }, []);

  const filteredHeads = selectedGroup ? allHeads.filter(h => h.groupId === selectedGroup) : allHeads;

  // Build layer grid (0-11)
  const layerGrid = useMemo(() => {
    const grid = Array.from({ length: 12 }, () => Array(12).fill(null));
    allHeads.forEach(h => {
      if (h.layer < 12 && h.head < 12) {
        grid[h.layer][h.head] = h;
      }
    });
    return grid;
  }, [allHeads]);

  // Toggle head ablation
  const toggleAblation = (layer, head) => {
    const key = `${layer}.${head}`;
    setKnockedOutHeads(prev =>
      prev.includes(key) ? prev.filter(k => k !== key) : [...prev, key]
    );
    setKnoutResultNull();
  };

  const setKnoutResultNull = () => {
    setKnockoutResult(null);
    setError(null);
  };

  const clearAllAblations = () => {
    setKnockedOutHeads([]);
    setKnockoutResult(null);
    setError(null);
  };

  const handleAblationModeChange = (mode) => {
    setAblationMode(mode);
    setKnockoutResult(null);
    setError(null);
  };

  // Run Knockout Request
  const runKnockout = useCallback(async () => {
    if (knockedOutHeads.length === 0) return;
    setLoading(true);
    setError(null);
    try {
      const headsList = knockedOutHeads.map(k => {
        const [l, h] = k.split('.').map(Number);
        return [l, h];
      });

      const res = await fetch(`${BACKEND_URL}/api/knockout`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          prompt: prompt.trim(),
          heads_to_knockout: headsList,
          mode: ablationMode
        })
      });

      if (!res.ok) {
        const errData = await res.json().catch(() => ({}));
        throw new Error(errData.detail || `Knockout failed (${res.status})`);
      }

      const data = await res.json();
      setKnockoutResult(data);
    } catch (e) {
      setError(e.message || 'Failed to complete attention head knockout');
    } finally {
      setLoading(false);
    }
  }, [knockedOutHeads, prompt, ablationMode]);

  return (
    <section id="circuits" className="scroll-mt-section" data-testid="section-circuit-explorer" style={{ padding: '80px 0' }}>
      <div className="section-container">
        {/* Header */}
        <div className="mb-2"><span className="badge-teal">Circuit Analysis</span></div>
        <h2 style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 600, fontSize: 30, color: '#E8EEF8', marginBottom: 8 }}>
          The IOI Circuit in GPT-2 Small
        </h2>
        <p style={{ fontSize: 15, color: '#8A9BC4', maxWidth: 680, marginBottom: 32, lineHeight: 1.75 }}>
          Reproduced from Wang et al. (2022). 26 attention heads across 7 functional groups &mdash; the largest end-to-end circuit reverse-engineered in a language model at time of publication.
        </p>

        {/* Task Explainer */}
        <div className="research-card-accent mb-8">
          <h3 style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 600, fontSize: 18, color: '#E8EEF8', marginBottom: 8 }}>
            The Indirect Object Identification Task
          </h3>
          <p style={{ fontSize: 14, color: '#8A9BC4', lineHeight: 1.75 }}>
            Given: <code style={{ color: '#00D9C0', background: '#0A0D18', padding: '2px 6px', borderRadius: 4, fontSize: 13 }}>"When Mary and John went to the store, John gave a bottle of milk to"</code>
            <br />GPT-2 Small correctly predicts <strong style={{ color: '#00E676' }}>" Mary"</strong> (not " John"). A circuit of 26 attention heads collaborate: some track which name is the subject, some detect the duplicate, some suppress the subject, and Name Mover heads move the indirect object to the output.
          </p>
        </div>

        {/* Group Filters */}
        <div className="flex flex-wrap gap-2 mb-6">
          <button
            onClick={() => setSelectedGroup(null)}
            className="text-xs px-3 py-1.5 rounded-full transition-colors duration-200 cursor-pointer"
            style={{
              background: !selectedGroup ? 'rgba(0,217,192,0.15)' : 'transparent',
              border: `1px solid ${!selectedGroup ? 'rgba(0,217,192,0.4)' : '#2A3A58'}`,
              color: !selectedGroup ? '#00D9C0' : '#8A9BC4',
            }}
          >
            All Groups ({circuitData.meta.total_heads})
          </button>
          {circuitData.groups.map(g => (
            <button
              key={g.id}
              onClick={() => setSelectedGroup(selectedGroup === g.id ? null : g.id)}
              className="text-xs px-3 py-1.5 rounded-full transition-colors duration-200 cursor-pointer"
              style={{
                background: selectedGroup === g.id ? `${g.color}22` : 'transparent',
                border: `1px solid ${selectedGroup === g.id ? `${g.color}66` : '#2A3A58'}`,
                color: selectedGroup === g.id ? g.color : '#8A9BC4',
              }}
            >
              {g.name} ({g.heads.length})
            </button>
          ))}
        </div>

        <div className="grid grid-cols-1 lg:grid-cols-3 gap-6 mb-8">
          {/* Circuit Diagram - Layer Grid */}
          <div className="lg:col-span-2">
            <div className="research-card" style={{ overflowX: 'auto' }}>
              <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', marginBottom: 12 }}>
                <span style={{ fontSize: 12, color: '#4A5A7A' }}>Layer × Head Grid &mdash; click cells to ablate heads</span>
                {knockedOutHeads.length > 0 && (
                  <button onClick={clearAllAblations} className="flex items-center gap-1 text-[11px] text-[#FF5063] hover:underline bg-transparent border-0 cursor-pointer">
                    <RotateCcw size={10} /> Clear {knockedOutHeads.length} Ablated
                  </button>
                )}
              </div>
              <div style={{ display: 'grid', gridTemplateColumns: 'auto repeat(12, 1fr)', gap: 2, minWidth: 500 }}>
                {/* Header Row */}
                <div style={{ fontSize: 10, color: '#4A5A7A', padding: 4 }}></div>
                {Array.from({ length: 12 }, (_, i) => (
                  <div key={i} style={{ fontSize: 10, color: '#4A5A7A', textAlign: 'center', padding: 4 }}>H{i}</div>
                ))}
                {/* Data Rows */}
                {layerGrid.map((row, layer) => (
                  <React.Fragment key={layer}>
                    <div style={{ fontSize: 10, color: '#4A5A7A', padding: 4, display: 'flex', alignItems: 'center' }}>L{layer}</div>
                    {row.map((head, headIdx) => {
                      const isAblated = knockedOutHeads.includes(`${layer}.${headIdx}`);
                      const isCircuitMember = !!head;
                      
                      let bg = '#0A0D18';
                      if (isAblated) {
                        bg = '#FF5063';
                      } else if (isCircuitMember) {
                        bg = (selectedGroup && head.groupId !== selectedGroup ? `${head.groupColor}22` : `${head.groupColor}55`);
                      }
                      
                      return (
                        <div
                          key={headIdx}
                          className="heatmap-cell"
                          onMouseEnter={() => isCircuitMember && setHoveredHead(head)}
                          onMouseLeave={() => setHoveredHead(null)}
                          onClick={() => toggleAblation(layer, headIdx)}
                          style={{
                            width: '100%',
                            aspectRatio: '1',
                            background: bg,
                            border: isAblated 
                              ? '2px solid #FF5063' 
                              : (hoveredHead === head && head ? `2px solid ${head.groupColor}` : '1px solid #1E2B4533'),
                            borderRadius: 3,
                            cursor: 'pointer',
                            display: 'flex',
                            alignItems: 'center',
                            justifyContent: 'center',
                            fontSize: 8,
                            color: isAblated ? '#060810' : (isCircuitMember ? head.groupColor : 'transparent'),
                            fontWeight: isAblated ? '700' : 'normal',
                            fontFamily: "'JetBrains Mono', monospace",
                            transition: 'border-color 200ms, background 200ms, color 200ms',
                          }}
                        >
                          {isAblated ? '✖' : (isCircuitMember ? head.name.replace('L', '').replace('H', '.') : '')}
                        </div>
                      );
                    })}
                  </React.Fragment>
                ))}
              </div>
            </div>

            {/* Hovered Head Detail */}
            {hoveredHead && (
              <div className="research-card mt-3" style={{ borderLeft: `3px solid ${hoveredHead.groupColor}` }}>
                <div className="flex items-center gap-2 mb-2">
                  <span style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 600, fontSize: 16, color: hoveredHead.groupColor }}>{hoveredHead.name}</span>
                  <span className="badge-teal" style={{ background: `${hoveredHead.groupColor}15`, color: hoveredHead.groupColor, borderColor: `${hoveredHead.groupColor}40` }}>{hoveredHead.groupName}</span>
                </div>
                <div className="grid grid-cols-2 gap-4" style={{ fontSize: 13 }}>
                  <div><span style={{ color: '#4A5A7A' }}>Attends to:</span> <span style={{ color: '#E8EEF8' }}>{hoveredHead.attends_to}</span></div>
                  <div><span style={{ color: '#4A5A7A' }}>Score:</span> <span style={{ color: hoveredHead.groupColor }}>{hoveredHead.score}</span></div>
                  <div><span style={{ color: '#4A5A7A' }}>Logit contribution:</span> <span style={{ color: hoveredHead.logit_contribution > 0 ? '#00E676' : '#FF5063' }}>{hoveredHead.logit_contribution > 0 ? '+' : ''}{hoveredHead.logit_contribution}</span></div>
                  <div><span style={{ color: '#4A5A7A' }}>Composition:</span> <span style={{ color: '#8A9BC4' }}>{hoveredHead.composition}</span></div>
                </div>
              </div>
            )}
          </div>

          {/* Group Descriptions */}
          <div className="space-y-3">
            {(selectedGroup ? circuitData.groups.filter(g => g.id === selectedGroup) : circuitData.groups).map(g => (
              <div key={g.id} className="research-card" style={{ borderLeft: `3px solid ${g.color}`, padding: '14px 16px' }}>
                <div className="flex items-center gap-2 mb-1">
                  <div style={{ width: 8, height: 8, borderRadius: '50%', background: g.color }} />
                  <span style={{ fontSize: 13, fontWeight: 600, color: g.color }}>{g.name}</span>
                  <span style={{ fontSize: 11, color: '#4A5A7A' }}>L{g.layers}</span>
                </div>
                <p style={{ fontSize: 12, color: '#8A9BC4', lineHeight: 1.6 }}>{g.description}</p>
              </div>
            ))}
          </div>
        </div>

        {/* Live Knockout Intervention Deck */}
        <div className="research-card mb-8">
          <div style={{ display: 'flex', justifyContent: 'space-between', flexWrap: 'wrap', alignItems: 'center', gap: 12, marginBottom: 16 }}>
            <h3 style={{ fontFamily: "'Space Grotesk', sans-serif", fontSize: 16, color: '#E8EEF8', fontWeight: 600, display: 'flex', alignItems: 'center', gap: 8 }}>
              <Zap size={18} style={{ color: '#FF5063' }} /> Live Head Knockout Intervention Lab
            </h3>
            <span className="badge-green">
              GPT-2 Active
            </span>
          </div>

          <p style={{ fontSize: 13, color: '#8A9BC4', marginBottom: 20, lineHeight: 1.6 }}>
            Select one or more attention heads on the grid above (such as the critical S-Inhibition heads <strong>L7.3</strong> or Induction heads <strong>L5.1</strong>) to ablate their attention patterns. Observe the downstream effect on logit difference necessity.
          </p>

          <div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
            <div className="lg:col-span-2 flex flex-col gap-4">
              {/* Ablation Mode Segmented Controller */}
              <div style={{ display: 'flex', alignItems: 'center', flexWrap: 'wrap', gap: '12px', background: '#080B12', border: '1px solid #1E2B45', borderRadius: '8px', padding: '8px 12px' }}>
                <span style={{ fontSize: '11px', fontWeight: 600, color: '#8A9BC4', fontFamily: "'Space Grotesk', sans-serif", textTransform: 'uppercase', letterSpacing: '0.05em' }}>Ablation Mode:</span>
                <div style={{ display: 'flex', gap: '4px', background: '#121729', borderRadius: '6px', padding: '2px', border: '1px solid #1E2B4533' }}>
                  <button
                    onClick={() => handleAblationModeChange('zero')}
                    style={{
                      background: ablationMode === 'zero' ? '#FF5063' : 'transparent',
                      color: ablationMode === 'zero' ? '#060810' : '#8A9BC4',
                      border: 'none',
                      borderRadius: '4px',
                      padding: '4px 10px',
                      fontSize: '11px',
                      fontWeight: 700,
                      cursor: 'pointer',
                      transition: 'all 200ms ease'
                    }}
                  >
                    Zero Ablation
                  </button>
                  <button
                    onClick={() => handleAblationModeChange('mean')}
                    style={{
                      background: ablationMode === 'mean' ? '#FF5063' : 'transparent',
                      color: ablationMode === 'mean' ? '#060810' : '#8A9BC4',
                      border: 'none',
                      borderRadius: '4px',
                      padding: '4px 10px',
                      fontSize: '11px',
                      fontWeight: 700,
                      cursor: 'pointer',
                      transition: 'all 200ms ease'
                    }}
                  >
                    Mean Ablation
                  </button>
                </div>
                <div style={{ fontSize: '11px', color: '#4A5A7A', flex: 1, minWidth: '180px', lineHeight: 1.3 }}>
                  {ablationMode === 'zero' 
                    ? "Zeroes attention weights. Can trigger out-of-distribution (OOD) downstream noise."
                    : "Averages attention patterns across tokens. Preserves natural activation bounds."
                  }
                </div>
              </div>

              <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: '10px 14px',
                    color: '#E8EEF8',
                    fontSize: 13,
                    fontFamily: "'Space Grotesk', sans-serif",
                    outline: 'none',
                  }}
                />
                <button
                  onClick={runKnockout}
                  disabled={loading || knockedOutHeads.length === 0}
                  className="btn-primary flex items-center gap-2"
                  style={{
                    padding: '10px 20px',
                    background: knockedOutHeads.length > 0 ? '#FF5063' : '#1E2B45',
                    borderColor: knockedOutHeads.length > 0 ? '#FF5063' : '#1E2B45',
                    color: knockedOutHeads.length > 0 ? '#060810' : '#4A5A7A',
                    fontWeight: 650
                  }}
                >
                  {loading ? <Loader2 className="animate-spin" size={16} /> : <Play size={16} />}
                  Ablate ({knockedOutHeads.length})
                </button>
              </div>

              {knockedOutHeads.length === 0 && (
                <div style={{ display: 'flex', alignItems: 'center', gap: 6, fontSize: 12, color: '#4A5A7A', fontFamily: "'JetBrains Mono', monospace" }}>
                  <Info size={13} /> Grid tip: Click any square on the grid above to load it into the knockout deck.
                </div>
              )}

              {error && (
                <div style={{ fontSize: 13, color: '#FF5063' }}>
                  <strong>Error:</strong> {error}
                </div>
              )}

              {knockoutResult && (
                <div className="research-card-accent" style={{ borderLeftColor: '#FF5063', marginTop: 4 }}>
                  <h4 style={{ fontSize: 13, color: '#E8EEF8', fontWeight: 600, marginBottom: 6 }}>Intervention Verdict: Head necessity is {knockoutResult.verdict.toUpperCase()}</h4>
                  <p style={{ fontSize: 13, color: '#8A9BC4', lineHeight: 1.6 }}>{knockoutResult.interpretation}</p>
                </div>
              )}
            </div>

            <div className="lg:col-span-1">
              <div style={{ background: '#080B12', border: '1px solid #1E2B4533', borderRadius: 8, padding: 16 }}>
                <div style={{ fontSize: 11, color: '#4A5A7A', marginBottom: 12, textTransform: 'uppercase', fontFamily: "'JetBrains Mono', monospace", fontWeight: 600 }}>
                  Ablation Performance Metrics
                </div>

                <div className="space-y-4">
                  <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
                    <span style={{ fontSize: 12, color: '#8A9BC4' }}>Baseline Logit Diff</span>
                    <span style={{ fontSize: 14, color: '#E8EEF8', fontWeight: 600, fontFamily: "'JetBrains Mono', monospace" }}>
                      {knockoutResult ? knockoutResult.baseline_logit_diff.toFixed(2) : '3.84'}
                    </span>
                  </div>

                  <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
                    <span style={{ fontSize: 12, color: '#8A9BC4' }}>Ablated Logit Diff</span>
                    <span style={{
                      fontSize: 14,
                      color: knockoutResult ? (knockoutResult.performance_retained_pct < 50 ? '#FF5063' : '#E8EEF8') : '#8A9BC4',
                      fontWeight: 600,
                      fontFamily: "'JetBrains Mono', monospace"
                    }}>
                      {knockoutResult ? knockoutResult.knocked_logit_diff.toFixed(2) : '--'}
                    </span>
                  </div>

                  <div style={{ height: 1, background: '#1E2B4533' }} />

                  <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
                    <span style={{ fontSize: 12, color: '#8A9BC4' }}>Performance Retained</span>
                    <span style={{
                      fontSize: 16,
                      color: knockoutResult ? (knockoutResult.performance_retained_pct < 50 ? '#FF5063' : '#00E676') : '#8A9BC4',
                      fontWeight: 700,
                      fontFamily: "'JetBrains Mono', monospace"
                    }}>
                      {knockoutResult ? `${knockoutResult.performance_retained_pct}%` : '--'}
                    </span>
                  </div>
                </div>
              </div>
            </div>
          </div>
        </div>

        {/* Faithfulness Metrics */}
        <div className="grid grid-cols-1 md:grid-cols-3 gap-4">
          <div className="stat-card" style={{ borderTop: '2px solid #00E676' }}>
            <div style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 700, fontSize: 40, color: '#00E676' }}>0.87</div>
            <div style={{ fontSize: 14, color: '#E8EEF8', fontWeight: 500, marginTop: 4 }}>Faithfulness</div>
            <p style={{ fontSize: 12, color: '#4A5A7A', marginTop: 4, lineHeight: 1.5 }}>
              The circuit alone recovers 87% of GPT-2 Small's logit difference on the IOI task. Full model F(M) = 3.56, circuit gap = 0.46.
            </p>
          </div>
          <div className="stat-card" style={{ borderTop: '2px solid #00D9C0' }}>
            <div style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 700, fontSize: 40, color: '#00D9C0' }}>0.91</div>
            <div style={{ fontSize: 14, color: '#E8EEF8', fontWeight: 500, marginTop: 4 }}>Completeness</div>
            <p style={{ fontSize: 12, color: '#4A5A7A', marginTop: 4, lineHeight: 1.5 }}>
              Knocking out the 26 heads degrades performance by 91% of total effect. The circuit is nearly complete.
            </p>
          </div>
          <div className="stat-card" style={{ borderTop: '2px solid #4A9EFF' }}>
            <div style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 700, fontSize: 40, color: '#4A9EFF' }}>Verified</div>
            <div style={{ fontSize: 14, color: '#E8EEF8', fontWeight: 500, marginTop: 4 }}>Minimality</div>
            <p style={{ fontSize: 12, color: '#4A5A7A', marginTop: 4, lineHeight: 1.5 }}>
              Each head contributes meaningfully. Removing any single head reduces performance by at least 5% of logit diff.
            </p>
          </div>
        </div>
      </div>
    </section>
  );
};