frontend/src/components/FeatureSteering.js

import React, { useState, useEffect, useCallback } from 'react';
import { Sliders, Sparkles, Send, Loader2, Info, Globe, Code, Play, Activity } from 'lucide-react';

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

const PRESETS = [
  {
    index: 2048,
    name: 'French Language',
    label: 'French Language Detector',
    category: 'language',
    icon: Globe,
    color: '#9B59F5',
    desc: 'Injects French syntactic and lexical directions, causing the model to complete sentences in French.',
    defaultPrompt: 'When Mary and John went to the store, John',
    strength: 15.0,
  },
  {
    index: 847,
    name: 'Python Programming',
    label: 'Python Code Detector',
    category: 'code',
    icon: Code,
    color: '#00D9C0',
    desc: 'Redirects representations into Python programming code blocks, generating function calls or syntax.',
    defaultPrompt: 'When Mary and John went to the store, John',
    strength: 12.0,
  },
  {
    index: 1580,
    name: 'Chess Notation',
    label: 'Chess Notation Detector',
    category: 'chess',
    icon: Play,
    color: '#4A9EFF',
    desc: 'Biases hidden representations toward chess tournament layouts and move notations.',
    defaultPrompt: 'When Mary and John went to the store, John',
    strength: 18.0,
  },
  {
    index: 3102,
    name: 'Medical Terminology',
    label: 'Medical Terms Detector',
    category: 'medical',
    icon: Activity,
    color: '#FF5063',
    desc: 'Steers hidden states into clinical contexts, changing general words into medical formulations.',
    defaultPrompt: 'When Mary and John went to the store, John',
    strength: 14.0,
  },
];

export const FeatureSteering = () => {
  const [activePreset, setActivePreset] = useState(PRESETS[0]);
  const [featureIndex, setFeatureIndex] = useState(PRESETS[0].index);
  const [steeringStrength, setSteeringStrength] = useState(PRESETS[0].strength);
  const [prompt, setPrompt] = useState(PRESETS[0].defaultPrompt);
  
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState(null);
  const [steerResult, setSteerResult] = 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 handleSelectPreset = (preset) => {
    setActivePreset(preset);
    setFeatureIndex(preset.index);
    setSteeringStrength(preset.strength);
    setPrompt(preset.defaultPrompt);
    setSteerResult(null);
  };

  const handleSteer = useCallback(async () => {
    if (!prompt.trim()) return;

    setLoading(true);
    setError(null);
    try {
      const res = await fetch(`${BACKEND_URL}/api/steer`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          prompt: prompt.trim(),
          feature_index: parseInt(featureIndex, 10),
          steering_strength: parseFloat(steeringStrength),
          layer: 6,
        }),
      });

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

      const data = await res.json();
      setSteerResult(data);
    } catch (e) {
      setError(e.message || 'Failed to execute feature steering');
    } finally {
      setLoading(false);
    }
  }, [prompt, featureIndex, steeringStrength]);

  return (
    <section id="steering" className="scroll-mt-section" data-testid="section-feature-steering" style={{ padding: '80px 0', borderTop: '1px solid #1E2B4533' }}>
      <div className="section-container">
        {/* Header */}
        <div className="mb-2">
          <span className="badge-violet">Phase 2: Feature Steering</span>
          <span style={{ marginLeft: 8 }} className="badge-green">
            GPT-2 + SAELens Active
          </span>
        </div>

        <h2 style={{ fontFamily: "'Space Grotesk', sans-serif", fontWeight: 600, fontSize: 30, color: '#E8EEF8', marginBottom: 8 }}>
          SAE Feature Steering Lab
        </h2>
        <p style={{ fontSize: 15, color: '#8A9BC4', maxWidth: 680, marginBottom: 32, lineHeight: 1.75 }}>
          Intervene directly in the model's computation by adding or subtracting pre-trained Sparse Autoencoder directions ({"$W_{dec}$"} decoders) scaled by strength parameters directly during generation.
        </p>

        {/* Preset Selector Grid */}
        <div className="grid grid-cols-1 md:grid-cols-4 gap-4 mb-8">
          {PRESETS.map((preset) => {
            const Icon = preset.icon;
            const isSelected = activePreset.index === preset.index;
            return (
              <div
                key={preset.index}
                onClick={() => handleSelectPreset(preset)}
                className="research-card transition-all duration-300"
                style={{
                  borderLeft: `3px solid ${preset.color}`,
                  cursor: 'pointer',
                  borderColor: isSelected ? preset.color : '#1E2B45',
                  boxShadow: isSelected ? `0 0 16px ${preset.color}15` : 'none',
                  background: isSelected ? '#0d111d' : '#080B12',
                }}
              >
                <div className="flex items-center gap-3 mb-2">
                  <div style={{
                    padding: 8,
                    borderRadius: 8,
                    background: `${preset.color}15`,
                    color: preset.color,
                  }}>
                    <Icon size={18} />
                  </div>
                  <div>
                    <h3 style={{ fontSize: 14, fontWeight: 600, color: '#E8EEF8' }}>{preset.name}</h3>
                    <span style={{ fontSize: 11, color: '#4A5A7A', fontFamily: "'JetBrains Mono', monospace" }}>Idx: {preset.index}</span>
                  </div>
                </div>
                <p style={{ fontSize: 12, color: '#8A9BC4', lineHeight: 1.5, marginTop: 4 }}>{preset.desc}</p>
              </div>
            );
          })}
        </div>

        <div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
          {/* Controls Panel */}
          <div className="lg:col-span-1">
            <div className="research-card">
              <h3 style={{ fontFamily: "'Space Grotesk', sans-serif", fontSize: 16, color: '#E8EEF8', fontWeight: 600, marginBottom: 20, display: 'flex', alignItems: 'center', gap: 8 }}>
                <Sliders size={18} style={{ color: activePreset.color }} /> Intervention Parameters
              </h3>

              {/* Feature Index */}
              <div className="mb-6">
                <label style={{ fontSize: 12, color: '#8A9BC4', display: 'block', marginBottom: 8, fontWeight: 500 }}>
                  Feature Latent Index (0 - 24575)
                </label>
                <input
                  type="number"
                  min="0"
                  max="24575"
                  value={featureIndex}
                  onChange={(e) => {
                    const idx = parseInt(e.target.value, 10) || 0;
                    setFeatureIndex(idx);
                    const matching = PRESETS.find(p => p.index === idx);
                    if (matching) {
                      setActivePreset(matching);
                    } else {
                      setActivePreset({
                        index: idx,
                        name: `Custom Latent ${idx}`,
                        label: `Feature ${idx}`,
                        color: '#9B59F5',
                        desc: 'Intervene with a custom dictionary element mapped from Neuronpedia API.',
                        icon: Sparkles,
                        defaultPrompt: prompt,
                        strength: steeringStrength
                      });
                    }
                  }}
                  disabled={loading}
                  style={{
                    width: '100%',
                    background: '#080B12',
                    border: '1px solid #1E2B45',
                    borderRadius: 6,
                    padding: '8px 12px',
                    color: '#E8EEF8',
                    fontFamily: "'JetBrains Mono', monospace",
                    fontSize: 13,
                    outline: 'none',
                  }}
                />
              </div>

              {/* Steering Strength */}
              <div className="mb-6">
                <div style={{ display: 'flex', justifyContent: 'space-between', marginBottom: 8 }}>
                  <label style={{ fontSize: 12, color: '#8A9BC4', fontWeight: 500 }}>
                    Steering Coefficient ($c$)
                  </label>
                  <span style={{ fontSize: 12, color: activePreset.color, fontWeight: 600, fontFamily: "'JetBrains Mono', monospace" }}>
                    {steeringStrength > 0 ? `+${steeringStrength}` : steeringStrength}
                  </span>
                </div>
                
                <input
                  type="range"
                  min="-30"
                  max="30"
                  step="0.5"
                  value={steeringStrength}
                  onChange={(e) => setSteeringStrength(parseFloat(e.target.value))}
                  disabled={loading}
                  style={{
                    width: '100%',
                    accentColor: activePreset.color,
                    cursor: 'pointer',
                  }}
                />

                <div style={{ display: 'flex', justifyContent: 'space-between', fontSize: 10, color: '#4A5A7A', marginTop: 6, fontFamily: "'JetBrains Mono', monospace" }}>
                  <span>-30.0 (Ablate)</span>
                  <span>0.0 (None)</span>
                  <span>+30.0 (Steer)</span>
                </div>
              </div>

              {/* Prompt Input */}
              <div className="mb-6">
                <label style={{ fontSize: 12, color: '#8A9BC4', display: 'block', marginBottom: 8, fontWeight: 500 }}>
                  Prompt Sequence
                </label>
                <textarea
                  rows="3"
                  value={prompt}
                  onChange={(e) => setPrompt(e.target.value)}
                  disabled={loading}
                  placeholder="Input sentence to generate..."
                  style={{
                    width: '100%',
                    background: '#080B12',
                    border: '1px solid #1E2B45',
                    borderRadius: 6,
                    padding: '10px 12px',
                    color: '#E8EEF8',
                    fontSize: 13,
                    fontFamily: "'Space Grotesk', sans-serif",
                    lineHeight: 1.5,
                    resize: 'none',
                    outline: 'none',
                  }}
                />
              </div>

              <button
                onClick={handleSteer}
                disabled={loading || !prompt.trim()}
                className="btn-primary w-full flex items-center justify-center gap-2"
                style={{
                  background: activePreset.color,
                  borderColor: activePreset.color,
                  color: '#060810',
                  fontWeight: 650,
                  padding: '10px 16px',
                }}
              >
                {loading ? <Loader2 className="animate-spin" size={16} /> : <Send size={16} />}
                Run Steering Hook
              </button>
              
              {error && (
                <div style={{ marginTop: 12, color: '#FF5063', fontSize: 12 }}>
                  <strong>Error:</strong> {error}
                </div>
              )}
            </div>
          </div>

          {/* Results Comparison View */}
          <div className="lg:col-span-2 flex flex-col gap-4">
            <div className="research-card flex-1 flex flex-col">
              <h3 style={{ fontFamily: "'Space Grotesk', sans-serif", fontSize: 16, color: '#E8EEF8', fontWeight: 600, marginBottom: 16, display: 'flex', alignItems: 'center', gap: 8 }}>
                <Sparkles size={18} style={{ color: activePreset.color }} /> Steered Completion Trace
              </h3>

              <div style={{ display: 'grid', gridTemplateColumns: '1fr', gap: 16, flex: 1 }}>
                {/* Original Completion */}
                <div style={{
                  background: '#080B12',
                  border: '1px solid #1E2B4533',
                  borderRadius: 8,
                  padding: 16,
                  display: 'flex',
                  flexDirection: 'column',
                }}>
                  <div style={{ fontSize: 11, color: '#4A5A7A', marginBottom: 8, fontWeight: 600, textTransform: 'uppercase', fontFamily: "'JetBrains Mono', monospace" }}>
                    Standard Completion (No steering)
                  </div>
                  <div style={{ fontSize: 14, color: '#8A9BC4', lineHeight: 1.6, flex: 1 }}>
                    <span style={{ color: '#E8EEF8' }}>{prompt}</span>{' '}
                    <span style={{ color: '#4A9EFF', fontWeight: 500 }}>
                      {steerResult ? steerResult.original_completion : '...'}
                    </span>
                  </div>
                </div>

                {/* Steered Completion */}
                <div style={{
                  background: '#0d111d',
                  border: `1px dashed ${activePreset.color}40`,
                  borderRadius: 8,
                  padding: 16,
                  display: 'flex',
                  flexDirection: 'column',
                  position: 'relative',
                  boxShadow: steerResult ? `inset 0 0 20px ${activePreset.color}05` : 'none',
                }}>
                  <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', marginBottom: 8 }}>
                    <div style={{ fontSize: 11, color: activePreset.color, fontWeight: 600, textTransform: 'uppercase', fontFamily: "'JetBrains Mono', monospace" }}>
                      Intervened Completion ({"$x \\leftarrow x + c \\cdot W_{dec}$"})
                    </div>
                    {steerResult && (
                      <span style={{ fontSize: 10, color: '#00D9C0', fontFamily: "'JetBrains Mono', monospace", background: '#00D9C010', padding: '2px 6px', borderRadius: 4 }}>
                        Steer successful
                      </span>
                    )}
                  </div>
                  <div style={{ fontSize: 14, color: '#8A9BC4', lineHeight: 1.6, flex: 1 }}>
                    <span style={{ color: '#E8EEF8' }}>{prompt}</span>{' '}
                    <span style={{ color: activePreset.color, fontWeight: 600 }}>
                      {steerResult ? steerResult.steered_completion : '...'}
                    </span>
                  </div>
                </div>
              </div>
            </div>

            {/* Explainer / Math Card */}
            <div className="research-card-accent" style={{ borderLeftColor: activePreset.color }}>
              <div style={{ display: 'flex', gap: 10, alignItems: 'flex-start' }}>
                <Info size={16} style={{ color: activePreset.color, marginTop: 2, flexShrink: 0 }} />
                <div>
                  <h4 style={{ fontSize: 13, color: '#E8EEF8', fontWeight: 600, marginBottom: 4 }}>
                    Mathematical Intervention Details
                  </h4>
                  <p style={{ fontSize: 12, color: '#8A9BC4', lineHeight: 1.6 }}>
                    This steering interface hooks into <code>blocks.6.hook_resid_post</code>. During the autoregressive generation loop, at each step, the residual stream tensor {"$x_6$"} is modified: {"$x_6 \\leftarrow x_6 + c \\cdot \\vec{v}$"}, where {"$\\vec{v}$"} is the normalized decoder dictionary column vector {"$W_{dec}[i]$"} associated with feature <strong>{steerResult ? steerResult.feature_label : activePreset.label}</strong> (index {featureIndex}). This injects semantic concept direction directly into GPT-2's core computational loop.
                  </p>
                </div>
              </div>
            </div>
          </div>
        </div>
      </div>
    </section>
  );
};