const { useState, useMemo, useCallback } = React;

// Project 3D point to 2D screen
const project = (point, rotX, rotY, scale = 100, offset = [200, 200]) => {
  let [x, y, z] = point;
  const cosY = Math.cos(rotY), sinY = Math.sin(rotY);
  [x, z] = [x * cosY - z * sinY, x * sinY + z * cosY];
  const cosX = Math.cos(rotX), sinX = Math.sin(rotX);
  [y, z] = [y * cosX - z * sinX, y * sinX + z * cosX];
  return [offset[0] + x * scale, offset[1] - y * scale, z];
};

const relu = (x) => Math.max(0, x);

const neuronActivation = (point, w, b) => {
  const preActivation = w[0] * point[0] + w[1] * point[1] + b;
  return { pre: preActivation, post: relu(preActivation) };
};

// Dataset generators
const datasets = {
  xor: {
    name: 'XOR Pattern',
    description: 'Classic XOR: opposite corners share class',
    generate: () => {
      const points = [];
      const class0Centers = [[-0.7, 0.7], [0.7, -0.7]];
      const class1Centers = [[0.7, 0.7], [-0.7, -0.7]];
      
      class0Centers.forEach(center => {
        for (let i = 0; i < 5; i++) {
          const angle = (i / 5) * Math.PI * 2 + Math.random() * 0.5;
          const r = 0.1 + Math.random() * 0.12;
          points.push({ pos: [center[0] + Math.cos(angle) * r, center[1] + Math.sin(angle) * r], class: 0 });
        }
      });
      class1Centers.forEach(center => {
        for (let i = 0; i < 5; i++) {
          const angle = (i / 5) * Math.PI * 2 + Math.random() * 0.5;
          const r = 0.1 + Math.random() * 0.12;
          points.push({ pos: [center[0] + Math.cos(angle) * r, center[1] + Math.sin(angle) * r], class: 1 });
        }
      });
      return points;
    },
    recommendedNeurons: [
      { w: [1, 1], b: 0 },
      { w: [1, -1], b: 0 },
      { w: [0, 1], b: 0 }
    ]
  },
  
  circles: {
    name: 'Concentric Circles',
    description: 'Inner circle vs outer ring',
    generate: () => {
      const points = [];
      // Inner circle - class 0
      for (let i = 0; i < 12; i++) {
        const angle = (i / 12) * Math.PI * 2 + Math.random() * 0.3;
        const r = 0.2 + Math.random() * 0.15;
        points.push({ pos: [Math.cos(angle) * r, Math.sin(angle) * r], class: 0 });
      }
      // Outer ring - class 1
      for (let i = 0; i < 14; i++) {
        const angle = (i / 14) * Math.PI * 2 + Math.random() * 0.2;
        const r = 0.7 + Math.random() * 0.2;
        points.push({ pos: [Math.cos(angle) * r, Math.sin(angle) * r], class: 1 });
      }
      return points;
    },
    recommendedNeurons: [
      { w: [1, 0], b: 0.4 },
      { w: [-1, 0], b: 0.4 },
      { w: [0, 1], b: 0.4 }
    ]
  },
  
  spiral: {
    name: 'Two Spirals',
    description: 'Interleaved spiral arms',
    generate: () => {
      const points = [];
      const turns = 1.5;
      for (let i = 0; i < 15; i++) {
        const t = i / 15;
        const angle = t * turns * Math.PI * 2;
        const r = 0.2 + t * 0.6;
        const noise = () => (Math.random() - 0.5) * 0.1;
        // Spiral 1
        points.push({ pos: [Math.cos(angle) * r + noise(), Math.sin(angle) * r + noise()], class: 0 });
        // Spiral 2 (offset by PI)
        points.push({ pos: [Math.cos(angle + Math.PI) * r + noise(), Math.sin(angle + Math.PI) * r + noise()], class: 1 });
      }
      return points;
    },
    recommendedNeurons: [
      { w: [1, 0.5], b: 0.2 },
      { w: [-0.5, 1], b: 0.2 },
      { w: [1, -1], b: 0 }
    ]
  },
  
  moons: {
    name: 'Two Moons',
    description: 'Interlocking crescent shapes',
    generate: () => {
      const points = [];
      // Upper moon - class 0
      for (let i = 0; i < 12; i++) {
        const angle = (i / 12) * Math.PI;
        const r = 0.5 + (Math.random() - 0.5) * 0.15;
        points.push({ pos: [Math.cos(angle) * r, Math.sin(angle) * r], class: 0 });
      }
      // Lower moon - class 1 (shifted and flipped)
      for (let i = 0; i < 12; i++) {
        const angle = Math.PI + (i / 12) * Math.PI;
        const r = 0.5 + (Math.random() - 0.5) * 0.15;
        points.push({ pos: [Math.cos(angle) * r + 0.5, Math.sin(angle) * r + 0.3], class: 1 });
      }
      return points;
    },
    recommendedNeurons: [
      { w: [1, 1], b: 0.3 },
      { w: [1, -1], b: 0.3 },
      { w: [-1, 0], b: 0.3 }
    ]
  },
  
  checkerboard: {
    name: 'Checkerboard',
    description: '2x2 alternating grid pattern',
    generate: () => {
      const points = [];
      const centers = [
        { pos: [-0.5, 0.5], class: 0 },
        { pos: [0.5, 0.5], class: 1 },
        { pos: [-0.5, -0.5], class: 1 },
        { pos: [0.5, -0.5], class: 0 }
      ];
      centers.forEach(({ pos, class: cls }) => {
        for (let i = 0; i < 5; i++) {
          const noise = () => (Math.random() - 0.5) * 0.3;
          points.push({ pos: [pos[0] + noise(), pos[1] + noise()], class: cls });
        }
      });
      return points;
    },
    recommendedNeurons: [
      { w: [1, 0], b: 0 },
      { w: [0, 1], b: 0 },
      { w: [1, 1], b: 0 }
    ]
  },
  
  stripe: {
    name: 'Middle Stripe',
    description: 'Class 1 forms a diagonal band',
    generate: () => {
      const points = [];
      for (let i = 0; i < 30; i++) {
        const x = (Math.random() - 0.5) * 2;
        const y = (Math.random() - 0.5) * 2;
        const inStripe = Math.abs(x + y) < 0.5;
        points.push({ pos: [x, y], class: inStripe ? 1 : 0 });
      }
      return points;
    },
    recommendedNeurons: [
      { w: [1, 1], b: 0.5 },
      { w: [-1, -1], b: 0.5 },
      { w: [1, -1], b: 0 }
    ]
  }
};

function NNGeometryVisualizer() {
  const [foldRotX, setFoldRotX] = useState(0.4);
  const [foldRotY, setFoldRotY] = useState(-0.3);
  const [hiddenRotX, setHiddenRotX] = useState(0.5);
  const [hiddenRotY, setHiddenRotY] = useState(0.4);
  const [dragging, setDragging] = useState(null);
  const [lastMouse, setLastMouse] = useState([0, 0]);
  
  const [liftAmount, setLiftAmount] = useState(0);
  const [showGrid, setShowGrid] = useState(true);
  const [showDecisionPlane, setShowDecisionPlane] = useState(true);
  const [activeNeuron, setActiveNeuron] = useState(null);
  const [highlightedPoint, setHighlightedPoint] = useState(null);
  
  const [selectedDataset, setSelectedDataset] = useState('xor');
  const [neuronCount, setNeuronCount] = useState(3);
  const [neurons, setNeurons] = useState(datasets.xor.recommendedNeurons);
  const [dataKey, setDataKey] = useState(0); // For regenerating data
  
  // Output weights adapt to neuron count
  const outputWeights = useMemo(() => {
    if (neuronCount === 1) return [1];
    if (neuronCount === 2) return [1, -1];
    return [0.8, 0.8, 1.2];
  }, [neuronCount]);
  const outputBias = -0.5;
  
  // Generate data points
  const dataPoints = useMemo(() => {
    return datasets[selectedDataset].generate();
  }, [selectedDataset, dataKey]);
  
  // Active neurons (limited by count)
  const activeNeurons = useMemo(() => neurons.slice(0, neuronCount), [neurons, neuronCount]);
  
  // Compute data in all spaces
  const processedData = useMemo(() => {
    return dataPoints.map((d, idx) => {
      const activations = activeNeurons.map(n => neuronActivation(d.pos, n.w, n.b));
      const h = activations.map(a => a.post);
      
      // Pad to 3D for visualization
      while (h.length < 3) h.push(0);
      
      const foldedZ = (h[0] * 0.4 + h[1] * 0.4 + h[2] * 0.3) * liftAmount;
      const folded = [d.pos[0], d.pos[1], foldedZ];
      const hidden = [h[0], h[1], h[2]];
      
      const output = outputWeights.reduce((sum, w, i) => sum + w * (activations[i]?.post || 0), 0) + outputBias;
      
      return { ...d, idx, activations, h, folded, hidden, output };
    });
  }, [dataPoints, activeNeurons, liftAmount, outputWeights]);
  
  // Grid for folded view
  const gridPoints = useMemo(() => {
    const points = [];
    const res = 20, range = 1.4;
    for (let i = 0; i <= res; i++) {
      for (let j = 0; j <= res; j++) {
        const x = (i / res - 0.5) * 2 * range;
        const y = (j / res - 0.5) * 2 * range;
        const h = activeNeurons.map(n => relu(n.w[0] * x + n.w[1] * y + n.b));
        while (h.length < 3) h.push(0);
        const z = (h[0] * 0.4 + h[1] * 0.4 + h[2] * 0.3) * liftAmount;
        points.push({ pos: [x, y], lifted: [x, y, z], h });
      }
    }
    return points;
  }, [activeNeurons, liftAmount]);
  
  // Neuron lines
  const neuronLines = useMemo(() => {
    return activeNeurons.map(({ w, b }) => {
      const range = 2;
      if (Math.abs(w[1]) > Math.abs(w[0])) {
        return [[-range, (-w[0] * -range - b) / w[1]], [range, (-w[0] * range - b) / w[1]]];
      } else if (Math.abs(w[0]) > 0.01) {
        return [[(-w[1] * -range - b) / w[0], -range], [(-w[1] * range - b) / w[0], range]];
      }
      return [[-range, -b / (w[1] || 0.01)], [range, -b / (w[1] || 0.01)]];
    });
  }, [activeNeurons]);
  
  // Decision boundary in 2D
  const decisionBoundary2D = useMemo(() => {
    const points = [];
    const res = 35, range = 1.4;
    for (let i = 0; i < res; i++) {
      for (let j = 0; j < res; j++) {
        const corners = [
          [(i/res - 0.5) * 2 * range, (j/res - 0.5) * 2 * range],
          [((i+1)/res - 0.5) * 2 * range, (j/res - 0.5) * 2 * range],
          [((i+1)/res - 0.5) * 2 * range, ((j+1)/res - 0.5) * 2 * range],
          [(i/res - 0.5) * 2 * range, ((j+1)/res - 0.5) * 2 * range]
        ];
        const outputs = corners.map(p => {
          const h = activeNeurons.map(n => relu(n.w[0] * p[0] + n.w[1] * p[1] + n.b));
          return outputWeights.reduce((sum, w, i) => sum + w * (h[i] || 0), 0) + outputBias;
        });
        if (outputs.some(o => o > 0) && outputs.some(o => o < 0)) {
          const cx = (corners[0][0] + corners[2][0]) / 2;
          const cy = (corners[0][1] + corners[2][1]) / 2;
          const h = activeNeurons.map(n => relu(n.w[0] * cx + n.w[1] * cy + n.b));
          while (h.length < 3) h.push(0);
          const z = (h[0] * 0.4 + h[1] * 0.4 + h[2] * 0.3) * liftAmount;
          points.push([cx, cy, z]);
        }
      }
    }
    return points;
  }, [activeNeurons, liftAmount, outputWeights]);
  
  const handleMouseDown = useCallback((view) => (e) => {
    setDragging(view);
    setLastMouse([e.clientX, e.clientY]);
  }, []);
  
  const handleMouseMove = useCallback((e) => {
    if (!dragging) return;
    const dx = e.clientX - lastMouse[0];
    const dy = e.clientY - lastMouse[1];
    
    if (dragging === 'fold') {
      setFoldRotY(prev => prev + dx * 0.01);
      setFoldRotX(prev => Math.max(-Math.PI/2, Math.min(Math.PI/2, prev + dy * 0.01)));
    } else if (dragging === 'hidden') {
      setHiddenRotY(prev => prev + dx * 0.01);
      setHiddenRotX(prev => Math.max(-Math.PI/2, Math.min(Math.PI/2, prev + dy * 0.01)));
    }
    setLastMouse([e.clientX, e.clientY]);
  }, [dragging, lastMouse]);
  
  const handleMouseUp = useCallback(() => setDragging(null), []);
  
  const updateNeuron = (idx, field, subfield, value) => {
    setNeurons(prev => {
      const next = [...prev];
      if (subfield !== null) {
        next[idx] = { ...next[idx], [field]: [...next[idx][field]] };
        next[idx][field][subfield] = value;
      } else {
        next[idx] = { ...next[idx], [field]: value };
      }
      return next;
    });
  };
  
  const handleDatasetChange = (key) => {
    setSelectedDataset(key);
    setNeurons(datasets[key].recommendedNeurons);
    setNeuronCount(3);
    setDataKey(prev => prev + 1);
    setHighlightedPoint(null);
  };
  
  const handleNeuronCountChange = (count) => {
    setNeuronCount(count);
    setActiveNeuron(null);
  };
  
  const neuronColors = ['#ff6b6b', '#4ecdc4', '#ffd93d'];
  const gridRes = 21;
  
  const separationInfo = useMemo(() => {
    const correct = processedData.filter(d => (d.output > 0) === (d.class === 1)).length;
    return {
      accuracy: (correct / processedData.length * 100).toFixed(0),
      correct,
      total: processedData.length
    };
  }, [processedData]);

  return (
    <div 
      style={{
        minHeight: '100vh',
        background: 'linear-gradient(145deg, #0a0a10 0%, #12121f 50%, #0a0a10 100%)',
        color: '#e0e0e0',
        fontFamily: '"IBM Plex Sans", system-ui, sans-serif',
        padding: '20px',
        boxSizing: 'border-box'
      }}
      onMouseMove={handleMouseMove}
      onMouseUp={handleMouseUp}
      onMouseLeave={handleMouseUp}
    >
      <style>{`
        @import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Sans:wght@300;400;500;600&family=IBM+Plex+Mono:wght@400;500&display=swap');
        input[type="range"] {
          -webkit-appearance: none;
          background: rgba(255,255,255,0.1);
          border-radius: 4px;
          height: 5px;
        }
        input[type="range"]::-webkit-slider-thumb {
          -webkit-appearance: none;
          width: 12px;
          height: 12px;
          border-radius: 50%;
          background: #64ffda;
          cursor: pointer;
        }
        .dataset-btn {
          padding: 8px 12px;
          background: rgba(255,255,255,0.03);
          border: 1px solid rgba(255,255,255,0.08);
          border-radius: 6px;
          color: #888;
          font-size: 11px;
          cursor: pointer;
          transition: all 0.2s;
        }
        .dataset-btn:hover {
          background: rgba(255,255,255,0.06);
          color: #bbb;
        }
        .dataset-btn.active {
          background: rgba(100, 255, 218, 0.1);
          border-color: rgba(100, 255, 218, 0.3);
          color: #64ffda;
        }
        .neuron-count-btn {
          width: 36px;
          height: 36px;
          display: flex;
          align-items: center;
          justify-content: center;
          background: rgba(255,255,255,0.03);
          border: 1px solid rgba(255,255,255,0.1);
          border-radius: 8px;
          color: #888;
          font-size: 14px;
          font-weight: 500;
          cursor: pointer;
          transition: all 0.2s;
        }
        .neuron-count-btn:hover {
          background: rgba(255,255,255,0.06);
        }
        .neuron-count-btn.active {
          background: rgba(100, 255, 218, 0.15);
          border-color: #64ffda;
          color: #64ffda;
        }
      `}</style>
      
      <div style={{ maxWidth: '1500px', margin: '0 auto' }}>
        <div style={{ marginBottom: '16px' }}>
          <h1 style={{ fontSize: '22px', fontWeight: 600, marginBottom: '4px', color: '#fff' }}>
            Neural Network Geometry: How Hidden Layers Enable Separation
          </h1>
          <p style={{ color: '#666', fontSize: '12px' }}>
            Explore how neurons fold 2D input space to make non-linear patterns linearly separable
          </p>
        </div>
        
        {/* Dataset and config row */}
        <div style={{ 
          display: 'flex', 
          gap: '16px', 
          marginBottom: '16px',
          flexWrap: 'wrap',
          alignItems: 'stretch'
        }}>
          {/* Dataset selector */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            flex: '1 1 auto'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#666', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '10px'
            }}>
              Dataset — Choose a non-linearly-separable pattern
            </div>
            <div style={{ display: 'flex', gap: '8px', flexWrap: 'wrap' }}>
              {Object.entries(datasets).map(([key, { name }]) => (
                <button
                  key={key}
                  className={`dataset-btn ${selectedDataset === key ? 'active' : ''}`}
                  onClick={() => handleDatasetChange(key)}
                >
                  {name}
                </button>
              ))}
              <button
                className="dataset-btn"
                onClick={() => setDataKey(prev => prev + 1)}
                style={{ marginLeft: 'auto' }}
              >
                ↻ Regenerate
              </button>
            </div>
            <div style={{ fontSize: '11px', color: '#555', marginTop: '8px' }}>
              {datasets[selectedDataset].description}
            </div>
          </div>
          
          {/* Neuron count */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            minWidth: '180px'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#666', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '10px'
            }}>
              Hidden Neurons
            </div>
            <div style={{ display: 'flex', gap: '8px' }}>
              {[1, 2, 3].map(n => (
                <button
                  key={n}
                  className={`neuron-count-btn ${neuronCount === n ? 'active' : ''}`}
                  onClick={() => handleNeuronCountChange(n)}
                >
                  {n}
                </button>
              ))}
            </div>
            <div style={{ fontSize: '11px', color: '#555', marginTop: '8px' }}>
              {neuronCount === 1 ? 'Single fold line' : 
               neuronCount === 2 ? 'Two fold lines → 2D hidden' : 
               'Three fold lines → 3D hidden'}
            </div>
          </div>
          
          {/* Lift control + accuracy */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            minWidth: '200px'
          }}>
            <div style={{ 
              display: 'flex', 
              justifyContent: 'space-between',
              marginBottom: '6px'
            }}>
              <span style={{ fontSize: '10px', color: '#666', textTransform: 'uppercase' }}>Lift Amount</span>
              <span style={{ fontSize: '11px', color: '#64ffda', fontFamily: '"IBM Plex Mono"' }}>
                {(liftAmount * 100).toFixed(0)}%
              </span>
            </div>
            <input
              type="range" min="0" max="1.5" step="0.02"
              value={liftAmount}
              onChange={(e) => setLiftAmount(parseFloat(e.target.value))}
              style={{ width: '100%' }}
            />
            <div style={{
              marginTop: '10px',
              padding: '8px',
              background: separationInfo.accuracy === '100' ? 'rgba(74, 222, 128, 0.1)' : 'rgba(255, 255, 255, 0.02)',
              borderRadius: '6px',
              display: 'flex',
              justifyContent: 'space-between',
              alignItems: 'center'
            }}>
              <span style={{ fontSize: '10px', color: '#888' }}>Accuracy</span>
              <span style={{ 
                fontSize: '16px', 
                fontWeight: 600,
                color: separationInfo.accuracy === '100' ? '#4ade80' : 
                       parseInt(separationInfo.accuracy) > 80 ? '#fbbf24' : '#fff'
              }}>
                {separationInfo.accuracy}%
              </span>
            </div>
          </div>
        </div>
        
        {/* Main visualization row */}
        <div style={{ display: 'flex', gap: '12px', flexWrap: 'wrap', marginBottom: '12px' }}>
          
          {/* 2D Input Space */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            flex: '0 0 240px'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#888', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '8px',
              display: 'flex',
              justifyContent: 'space-between',
              alignItems: 'center'
            }}>
              <span>① Input (x, y)</span>
              <span style={{ 
                color: '#ff6b6b', 
                fontSize: '9px',
                background: 'rgba(255, 107, 107, 0.1)',
                padding: '2px 5px',
                borderRadius: '3px'
              }}>
                Not Separable
              </span>
            </div>
            
            <svg width="215" height="215" style={{ 
              background: 'rgba(5, 5, 10, 0.8)',
              borderRadius: '6px'
            }}>
              {/* Grid */}
              {Array.from({ length: 7 }).map((_, i) => {
                const pos = 20 + i * 29;
                return (
                  <g key={i}>
                    <line x1={pos} y1={20} x2={pos} y2={195} stroke="rgba(255,255,255,0.03)" />
                    <line x1={20} y1={pos} x2={195} y2={pos} stroke="rgba(255,255,255,0.03)" />
                  </g>
                );
              })}
              
              {/* Neuron lines */}
              {neuronLines.map((line, i) => {
                const color = neuronColors[i];
                const toScreen = (p) => [107 + p[0] * 60, 107 - p[1] * 60];
                const p1 = toScreen(line[0]);
                const p2 = toScreen(line[1]);
                const isActive = activeNeuron === i;
                
                return (
                  <line key={`line-${i}`}
                    x1={p1[0]} y1={p1[1]} x2={p2[0]} y2={p2[1]}
                    stroke={color}
                    strokeWidth={isActive ? 2 : 1.5}
                    strokeDasharray={isActive ? 'none' : '3,2'}
                    opacity={isActive ? 1 : 0.5}
                  />
                );
              })}
              
              {/* Decision boundary */}
              {decisionBoundary2D.map((p, i) => {
                const x = 107 + p[0] * 60;
                const y = 107 - p[1] * 60;
                return <circle key={`db-${i}`} cx={x} cy={y} r={1.5} fill="#a855f7" opacity={0.6} />;
              })}
              
              {/* Data points */}
              {processedData.map((d, i) => {
                const x = 107 + d.pos[0] * 60;
                const y = 107 - d.pos[1] * 60;
                const color = d.class === 0 ? '#64ffda' : '#ff6b9d';
                const isHighlighted = highlightedPoint === i;
                
                return (
                  <g key={`pt-${i}`}>
                    {isHighlighted && <circle cx={x} cy={y} r={10} fill="none" stroke={color} strokeWidth="1.5" opacity={0.4} />}
                    <circle
                      cx={x} cy={y} r={4}
                      fill={color}
                      stroke="#fff"
                      strokeWidth={isHighlighted ? 1 : 0.3}
                      style={{ cursor: 'pointer' }}
                      onMouseEnter={() => setHighlightedPoint(i)}
                      onMouseLeave={() => setHighlightedPoint(null)}
                    />
                  </g>
                );
              })}
            </svg>
            
            <div style={{ marginTop: '6px', fontSize: '10px', color: '#555' }}>
              <span style={{ color: '#64ffda' }}>●</span> Class 0 &nbsp;
              <span style={{ color: '#ff6b9d' }}>●</span> Class 1
            </div>
          </div>
          
          {/* Arrow */}
          <div style={{ display: 'flex', alignItems: 'center', padding: '0 4px' }}>
            <svg width="30" height="24">
              <defs>
                <marker id="arr" markerWidth="8" markerHeight="6" refX="7" refY="3" orient="auto">
                  <polygon points="0 0, 8 3, 0 6" fill="#64ffda" />
                </marker>
              </defs>
              <line x1="2" y1="12" x2="22" y2="12" stroke="#64ffda" strokeWidth="2" markerEnd="url(#arr)" />
            </svg>
          </div>
          
          {/* Hidden Space */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(100, 255, 218, 0.15)',
            flex: '1 1 300px'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#64ffda', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '8px',
              display: 'flex',
              justifyContent: 'space-between',
              alignItems: 'center'
            }}>
              <span>② Hidden Space ({neuronCount === 1 ? 'h₁' : neuronCount === 2 ? 'h₁, h₂' : 'h₁, h₂, h₃'})</span>
              <span style={{ 
                color: '#4ade80', 
                fontSize: '9px',
                background: 'rgba(74, 222, 128, 0.1)',
                padding: '2px 5px',
                borderRadius: '3px'
              }}>
                {neuronCount >= 2 ? 'Separable!' : 'Maybe...'}
              </span>
            </div>
            
            <svg
              width="100%"
              height="260"
              viewBox="0 0 350 260"
              style={{ 
                cursor: dragging === 'hidden' ? 'grabbing' : 'grab',
                background: 'radial-gradient(ellipse at center, #0f1015 0%, #080810 100%)',
                borderRadius: '6px'
              }}
              onMouseDown={handleMouseDown('hidden')}
            >
              {/* Decision plane */}
              {showDecisionPlane && neuronCount >= 2 && (
                <g>
                  {(() => {
                    const lines = [];
                    const range = 1.4;
                    const steps = 10;
                    
                    for (let i = 0; i <= steps; i++) {
                      for (let j = 0; j <= steps; j++) {
                        const h1 = (i / steps) * range;
                        const h2 = (j / steps) * range;
                        let h3 = 0;
                        
                        if (neuronCount === 3 && outputWeights[2]) {
                          h3 = -(outputWeights[0] * h1 + outputWeights[1] * h2 + outputBias) / outputWeights[2];
                          if (h3 < -0.1 || h3 > range * 1.2) continue;
                        } else if (neuronCount === 2) {
                          // For 2D, check if on the decision line
                          const val = outputWeights[0] * h1 + outputWeights[1] * h2 + outputBias;
                          if (Math.abs(val) > 0.15) continue;
                        }
                        
                        const proj = project([h1, h2, h3], hiddenRotX, hiddenRotY, 110, [175, 140]);
                        
                        // Draw connections to neighbors
                        const nextI = (i + 1) / steps * range;
                        const nextJ = (j + 1) / steps * range;
                        
                        if (i < steps) {
                          let h3Next = 0;
                          if (neuronCount === 3 && outputWeights[2]) {
                            h3Next = -(outputWeights[0] * nextI + outputWeights[1] * h2 + outputBias) / outputWeights[2];
                          }
                          if (h3Next >= -0.1 && h3Next <= range * 1.2) {
                            const projNext = project([nextI, h2, h3Next], hiddenRotX, hiddenRotY, 110, [175, 140]);
                            lines.push(
                              <line key={`ph-${i}-${j}`} x1={proj[0]} y1={proj[1]} x2={projNext[0]} y2={projNext[1]}
                                stroke="rgba(168, 85, 247, 0.25)" strokeWidth="1" />
                            );
                          }
                        }
                        if (j < steps) {
                          let h3Next = 0;
                          if (neuronCount === 3 && outputWeights[2]) {
                            h3Next = -(outputWeights[0] * h1 + outputWeights[1] * nextJ + outputBias) / outputWeights[2];
                          }
                          if (h3Next >= -0.1 && h3Next <= range * 1.2) {
                            const projNext = project([h1, nextJ, h3Next], hiddenRotX, hiddenRotY, 110, [175, 140]);
                            lines.push(
                              <line key={`pv-${i}-${j}`} x1={proj[0]} y1={proj[1]} x2={projNext[0]} y2={projNext[1]}
                                stroke="rgba(168, 85, 247, 0.25)" strokeWidth="1" />
                            );
                          }
                        }
                      }
                    }
                    return lines;
                  })()}
                </g>
              )}
              
              {/* Data points */}
              {processedData
                .map((d, i) => ({ ...d, proj: project(d.hidden, hiddenRotX, hiddenRotY, 110, [175, 140]) }))
                .sort((a, b) => a.proj[2] - b.proj[2])
                .map((d) => {
                  const color = d.class === 0 ? '#64ffda' : '#ff6b9d';
                  const isHighlighted = highlightedPoint === d.idx;
                  const correct = (d.output > 0) === (d.class === 1);
                  
                  return (
                    <g key={`hidden-${d.idx}`}>
                      {isHighlighted && (
                        <circle cx={d.proj[0]} cy={d.proj[1]} r={14}
                          fill="none" stroke={color} strokeWidth="1.5" opacity={0.4} />
                      )}
                      <circle
                        cx={d.proj[0]} cy={d.proj[1]} r={6}
                        fill={color}
                        stroke={correct ? '#fff' : '#ff0000'}
                        strokeWidth={isHighlighted ? 1.5 : 0.5}
                        style={{ cursor: 'pointer', filter: `drop-shadow(0 0 3px ${color})` }}
                        onMouseEnter={() => setHighlightedPoint(d.idx)}
                        onMouseLeave={() => setHighlightedPoint(null)}
                      />
                    </g>
                  );
              })}
              
              {/* Axes */}
              {[
                { dir: [1.2, 0, 0], color: neuronColors[0], label: 'h₁', show: neuronCount >= 1 },
                { dir: [0, 1.2, 0], color: neuronColors[1], label: 'h₂', show: neuronCount >= 2 },
                { dir: [0, 0, 1.2], color: neuronColors[2], label: 'h₃', show: neuronCount >= 3 }
              ].filter(a => a.show).map(({ dir, color, label }) => {
                const origin = project([0, 0, 0], hiddenRotX, hiddenRotY, 110, [175, 140]);
                const end = project(dir, hiddenRotX, hiddenRotY, 110, [175, 140]);
                return (
                  <g key={label}>
                    <line x1={origin[0]} y1={origin[1]} x2={end[0]} y2={end[1]}
                      stroke={color} strokeWidth={2} opacity={0.8} />
                    <text x={end[0] + 6} y={end[1] + 3} fill={color} fontSize="11" fontWeight="500">
                      {label}
                    </text>
                  </g>
                );
              })}
            </svg>
            
            <div style={{ marginTop: '6px', display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
              <span style={{ fontSize: '10px', color: '#666' }}>
                <span style={{ color: '#a855f7' }}>■</span> Decision {neuronCount === 2 ? 'line' : 'plane'}
              </span>
              <label style={{ fontSize: '10px', color: '#555', display: 'flex', alignItems: 'center', gap: '4px' }}>
                <input type="checkbox" checked={showDecisionPlane} onChange={(e) => setShowDecisionPlane(e.target.checked)}
                  style={{ accentColor: '#a855f7', width: '11px', height: '11px' }} />
                Show
              </label>
            </div>
          </div>
          
          {/* Arrow */}
          <div style={{ display: 'flex', alignItems: 'center', padding: '0 4px' }}>
            <svg width="30" height="24">
              <line x1="2" y1="12" x2="22" y2="12" stroke="#a855f7" strokeWidth="2" markerEnd="url(#arr)" />
            </svg>
          </div>
          
          {/* Folded View */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            flex: '1 1 300px'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#888', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '8px',
              display: 'flex',
              justifyContent: 'space-between'
            }}>
              <span>③ Folded View (x, y, lift)</span>
              <label style={{ fontSize: '9px', color: '#555', display: 'flex', alignItems: 'center', gap: '3px' }}>
                <input type="checkbox" checked={showGrid} onChange={(e) => setShowGrid(e.target.checked)}
                  style={{ accentColor: '#64ffda', width: '10px', height: '10px' }} />
                Grid
              </label>
            </div>
            
            <svg
              width="100%"
              height="260"
              viewBox="0 0 350 260"
              style={{ 
                cursor: dragging === 'fold' ? 'grabbing' : 'grab',
                background: 'radial-gradient(ellipse at center, #0f1015 0%, #080810 100%)',
                borderRadius: '6px'
              }}
              onMouseDown={handleMouseDown('fold')}
            >
              {/* Grid mesh */}
              {showGrid && gridPoints.map((point, idx) => {
                const i = idx % gridRes, j = Math.floor(idx / gridRes);
                if (i >= gridRes - 1 || j >= gridRes - 1) return null;
                
                const p1 = project(point.lifted, foldRotX, foldRotY, 85, [175, 145]);
                const p2 = project(gridPoints[idx + 1].lifted, foldRotX, foldRotY, 85, [175, 145]);
                const p3 = project(gridPoints[idx + gridRes].lifted, foldRotX, foldRotY, 85, [175, 145]);
                const alpha = 0.05 + Math.max(0, point.lifted[2]) * 0.06;
                
                return (
                  <g key={idx}>
                    <line x1={p1[0]} y1={p1[1]} x2={p2[0]} y2={p2[1]} 
                      stroke={`rgba(100, 255, 218, ${alpha})`} strokeWidth="0.5" />
                    <line x1={p1[0]} y1={p1[1]} x2={p3[0]} y2={p3[1]} 
                      stroke={`rgba(100, 255, 218, ${alpha})`} strokeWidth="0.5" />
                  </g>
                );
              })}
              
              {/* Fold lines */}
              {neuronLines.map((line, ni) => {
                const color = neuronColors[ni];
                const isActive = activeNeuron === ni;
                const segments = [];
                const steps = 20;
                
                for (let s = 0; s < steps; s++) {
                  const t1 = s / steps, t2 = (s + 1) / steps;
                  const p1 = [
                    line[0][0] + (line[1][0] - line[0][0]) * t1,
                    line[0][1] + (line[1][1] - line[0][1]) * t1
                  ];
                  const p2 = [
                    line[0][0] + (line[1][0] - line[0][0]) * t2,
                    line[0][1] + (line[1][1] - line[0][1]) * t2
                  ];
                  
                  const h1 = activeNeurons.map(n => relu(n.w[0] * p1[0] + n.w[1] * p1[1] + n.b));
                  const h2 = activeNeurons.map(n => relu(n.w[0] * p2[0] + n.w[1] * p2[1] + n.b));
                  while (h1.length < 3) h1.push(0);
                  while (h2.length < 3) h2.push(0);
                  const z1 = (h1[0] * 0.4 + h1[1] * 0.4 + h1[2] * 0.3) * liftAmount;
                  const z2 = (h2[0] * 0.4 + h2[1] * 0.4 + h2[2] * 0.3) * liftAmount;
                  
                  const proj1 = project([p1[0], p1[1], z1], foldRotX, foldRotY, 85, [175, 145]);
                  const proj2 = project([p2[0], p2[1], z2], foldRotX, foldRotY, 85, [175, 145]);
                  
                  segments.push(
                    <line key={`${ni}-${s}`}
                      x1={proj1[0]} y1={proj1[1]} x2={proj2[0]} y2={proj2[1]}
                      stroke={color} strokeWidth={isActive ? 2 : 1.5} opacity={isActive ? 1 : 0.5} />
                  );
                }
                return <g key={`fold-${ni}`}>{segments}</g>;
              })}
              
              {/* Decision boundary */}
              {decisionBoundary2D.map((p, i) => {
                const proj = project(p, foldRotX, foldRotY, 85, [175, 145]);
                return <circle key={`db-${i}`} cx={proj[0]} cy={proj[1]} r={1.5} fill="#a855f7" opacity={0.7} />;
              })}
              
              {/* Data points */}
              {processedData.map((d) => {
                const proj = project(d.folded, foldRotX, foldRotY, 85, [175, 145]);
                const color = d.class === 0 ? '#64ffda' : '#ff6b9d';
                const isHighlighted = highlightedPoint === d.idx;
                const correct = (d.output > 0) === (d.class === 1);
                
                return (
                  <g key={`fold-pt-${d.idx}`}>
                    {isHighlighted && <circle cx={proj[0]} cy={proj[1]} r={12} fill="none" stroke={color} strokeWidth="1.5" opacity={0.4} />}
                    <circle
                      cx={proj[0]} cy={proj[1]} r={5}
                      fill={color}
                      stroke={correct ? '#fff' : '#ff0000'}
                      strokeWidth={isHighlighted ? 1 : 0.4}
                      style={{ cursor: 'pointer' }}
                      onMouseEnter={() => setHighlightedPoint(d.idx)}
                      onMouseLeave={() => setHighlightedPoint(null)}
                    />
                  </g>
                );
              })}
              
              {/* Axes */}
              {[
                { dir: [1.3, 0, 0], color: '#555', label: 'x' },
                { dir: [0, 1.3, 0], color: '#555', label: 'y' },
                { dir: [0, 0, 1], color: '#ffd700', label: 'lift' }
              ].map(({ dir, color, label }) => {
                const origin = project([0, 0, 0], foldRotX, foldRotY, 85, [175, 145]);
                const end = project(dir, foldRotX, foldRotY, 85, [175, 145]);
                return (
                  <g key={label}>
                    <line x1={origin[0]} y1={origin[1]} x2={end[0]} y2={end[1]}
                      stroke={color} strokeWidth={1.5} opacity={0.5} />
                    <text x={end[0] + 5} y={end[1] + 3} fill={color} fontSize="9">{label}</text>
                  </g>
                );
              })}
            </svg>
            
            <div style={{ marginTop: '6px', fontSize: '10px', color: '#555' }}>
              Original coords + height from neuron outputs
            </div>
          </div>
        </div>
        
        {/* Neuron controls + point inspector */}
        <div style={{ display: 'flex', gap: '12px', flexWrap: 'wrap' }}>
          
          {/* Neuron controls */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            flex: '2 1 400px'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#888', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '10px'
            }}>
              Neuron Parameters — Adjust fold lines
            </div>
            
            <div style={{ display: 'flex', gap: '10px', flexWrap: 'wrap' }}>
              {neurons.slice(0, neuronCount).map((neuron, i) => (
                <div
                  key={i}
                  style={{
                    flex: '1 1 140px',
                    padding: '10px',
                    background: activeNeuron === i ? 'rgba(255,255,255,0.03)' : 'transparent',
                    borderRadius: '8px',
                    border: `1px solid ${activeNeuron === i ? neuronColors[i] : 'rgba(255,255,255,0.05)'}`,
                    cursor: 'pointer'
                  }}
                  onClick={() => setActiveNeuron(activeNeuron === i ? null : i)}
                >
                  <div style={{ display: 'flex', alignItems: 'center', gap: '6px', marginBottom: '8px' }}>
                    <div style={{
                      width: '8px', height: '8px', borderRadius: '50%',
                      background: neuronColors[i],
                      boxShadow: `0 0 6px ${neuronColors[i]}50`
                    }} />
                    <span style={{ fontSize: '11px', fontWeight: 500 }}>h{i + 1}</span>
                    <span style={{ fontSize: '9px', color: '#555', marginLeft: 'auto', fontFamily: '"IBM Plex Mono"' }}>
                      {neuron.w[0].toFixed(1)}x{neuron.w[1] >= 0 ? '+' : ''}{neuron.w[1].toFixed(1)}y{neuron.b >= 0 ? '+' : ''}{neuron.b.toFixed(1)}
                    </span>
                  </div>
                  
                  <div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr 1fr', gap: '6px' }}>
                    {['w₁', 'w₂', 'b'].map((label, li) => (
                      <div key={label}>
                        <div style={{ fontSize: '8px', color: '#555', marginBottom: '2px' }}>{label}</div>
                        <input 
                          type="range" 
                          min={li < 2 ? '-1.5' : '-1'} 
                          max={li < 2 ? '1.5' : '1'} 
                          step="0.1" 
                          value={li < 2 ? neuron.w[li] : neuron.b}
                          onChange={(e) => updateNeuron(i, li < 2 ? 'w' : 'b', li < 2 ? li : null, parseFloat(e.target.value))}
                          onClick={(e) => e.stopPropagation()}
                          style={{ width: '100%' }} 
                        />
                      </div>
                    ))}
                  </div>
                </div>
              ))}
              
              {/* Disabled neurons */}
              {neurons.slice(neuronCount).map((_, i) => (
                <div
                  key={`disabled-${i}`}
                  style={{
                    flex: '1 1 140px',
                    padding: '10px',
                    background: 'rgba(0,0,0,0.2)',
                    borderRadius: '8px',
                    border: '1px dashed rgba(255,255,255,0.1)',
                    opacity: 0.4
                  }}
                >
                  <div style={{ fontSize: '11px', color: '#555', textAlign: 'center' }}>
                    h{neuronCount + i + 1} disabled
                  </div>
                </div>
              ))}
            </div>
          </div>
          
          {/* Point inspector */}
          <div style={{
            background: 'rgba(18, 20, 30, 0.95)',
            borderRadius: '12px',
            padding: '14px',
            border: '1px solid rgba(255, 255, 255, 0.06)',
            flex: '1 1 220px',
            minWidth: '220px'
          }}>
            <div style={{ 
              fontSize: '10px', 
              color: '#888', 
              textTransform: 'uppercase',
              letterSpacing: '1px',
              marginBottom: '10px'
            }}>
              {highlightedPoint !== null ? 'Point Details' : 'Hover a Point'}
            </div>
            
            {highlightedPoint !== null ? (
              <div style={{ fontFamily: '"IBM Plex Mono"', fontSize: '10px' }}>
                <div style={{ 
                  display: 'flex', 
                  justifyContent: 'space-between',
                  marginBottom: '10px',
                  padding: '6px',
                  background: 'rgba(0,0,0,0.2)',
                  borderRadius: '4px'
                }}>
                  <span style={{ color: '#888' }}>True Class</span>
                  <span style={{ 
                    color: processedData[highlightedPoint].class === 0 ? '#64ffda' : '#ff6b9d',
                    fontWeight: 500
                  }}>
                    {processedData[highlightedPoint].class}
                  </span>
                </div>
                
                <div style={{ display: 'flex', flexDirection: 'column', gap: '4px' }}>
                  {activeNeurons.map((_, ni) => {
                    const act = processedData[highlightedPoint].activations[ni];
                    return (
                      <div key={ni} style={{ 
                        display: 'flex',
                        alignItems: 'center',
                        gap: '6px',
                        padding: '5px 6px',
                        background: 'rgba(0,0,0,0.2)',
                        borderRadius: '4px',
                        borderLeft: `2px solid ${neuronColors[ni]}`
                      }}>
                        <span style={{ color: '#666', width: '18px' }}>h{ni + 1}</span>
                        <span style={{ color: act.pre >= 0 ? '#4ade80' : '#666', width: '40px' }}>
                          {act.pre.toFixed(2)}
                        </span>
                        <span style={{ color: '#444' }}>→</span>
                        <span style={{ color: '#fff' }}>{act.post.toFixed(2)}</span>
                        {act.pre < 0 && (
                          <span style={{ color: '#ff6b6b', fontSize: '8px', marginLeft: 'auto' }}>zeroed</span>
                        )}
                      </div>
                    );
                  })}
                </div>
                
                <div style={{ 
                  marginTop: '8px',
                  padding: '6px',
                  background: (processedData[highlightedPoint].output > 0) === (processedData[highlightedPoint].class === 1)
                    ? 'rgba(74, 222, 128, 0.1)' 
                    : 'rgba(255, 107, 107, 0.1)',
                  borderRadius: '4px'
                }}>
                  <span style={{ color: '#888' }}>Output: </span>
                  <span style={{ color: processedData[highlightedPoint].output > 0 ? '#4ade80' : '#f87171' }}>
                    {processedData[highlightedPoint].output.toFixed(2)}
                  </span>
                  <span style={{ color: '#555' }}> → class {processedData[highlightedPoint].output > 0 ? 1 : 0}</span>
                </div>
              </div>
            ) : (
              <div style={{ color: '#444', fontSize: '11px', lineHeight: 1.5 }}>
                Hover points to see neuron activations and how ReLU transforms them.
              </div>
            )}
          </div>
        </div>
        
        {/* Compact explanation */}
        <div style={{
          marginTop: '12px',
          padding: '14px',
          background: 'rgba(18, 20, 30, 0.6)',
          borderRadius: '10px',
          border: '1px solid rgba(255, 255, 255, 0.04)',
          fontSize: '11px',
          color: '#777',
          display: 'flex',
          gap: '24px',
          flexWrap: 'wrap'
        }}>
          <div><strong style={{ color: '#64ffda' }}>Hidden Space</strong> — Each axis = one neuron's ReLU output (distance from line, positive side only)</div>
          <div><strong style={{ color: '#ffd700' }}>Folded View</strong> — Same transform shown as the input plane being creased along fold lines</div>
          <div><strong style={{ color: '#a855f7' }}>Decision Plane</strong> — A flat cut in hidden space → complex boundary in input space</div>
        </div>
      </div>
    </div>
  );
}

ReactDOM.render(<NNGeometryVisualizer />, document.getElementById('nn-app-root'));