import { Arrow, FILL, LoopFigure, type LoopStop } from "./LoopFigure";

const CX = [110, 280, 450, 620];
const INPUT_Y = [120, 200, 280, 360];
const HID_Y = [90, 165, 240, 315, 390];
const OUT_Y = [160, 240, 320];

const STOPS: LoopStop[] = [
  {
    x: CX[0],
    y: 200,
    ids: ["input"],
    inspect: "input",
    title: "Input layer",
    note: "A fixed-size vector x. No notion of sequence — one snapshot of features.",
  },
  {
    x: CX[1],
    y: 240,
    ids: ["hidden"],
    inspect: "hidden",
    title: "Hidden layer",
    note: "h = σ(W₁x + b₁). Width buys capacity; the same map is used for every example.",
  },
  {
    x: CX[2],
    y: 240,
    ids: ["hidden"],
    inspect: "hidden",
    title: "Deeper hidden layer",
    note: "Stacking nonlinear maps gives universal approximation — still one-shot, still no memory.",
  },
  {
    x: CX[3],
    y: 240,
    ids: ["output"],
    inspect: "output",
    title: "Output layer",
    note: "y is class scores or a regression target. Then the loop restarts at a new x.",
  },
  {
    x: CX[3],
    y: 455,
    ids: ["output"],
    inspect: "output",
    title: "Next example",
    note: "Each forward pass is independent. Nothing is carried to the next input.",
  },
  {
    x: CX[0],
    y: 455,
    ids: ["input"],
    inspect: "input",
    title: "Next example",
    note: "A new vector x enters the same stack of weights.",
  },
];

export function MlpLoop({ onSelect }: { onSelect: (id: string) => void }) {
  return (
    <LoopFigure
      stops={STOPS}
      loopMs={12000}
      viewBox="0 0 740 470"
      ariaLabel="Multilayer perceptron architecture"
      onSelect={onSelect}
    >
      {({ on, markerId }) => (
        <g>
          <rect width="740" height="470" fill="#f7f4ec" />
          <text x="370" y="32" textAnchor="middle" fontSize="16" fontWeight="700" fill="#222">
            Multilayer Perceptron
          </text>
          <text x="370" y="52" textAnchor="middle" fontSize="12" fill="#555">
            h = σ(Wx + b)
          </text>

          {INPUT_Y.map((y, i) =>
            HID_Y.map((hy, j) => (
              <line
                key={`i-${i}-${j}`}
                x1={CX[0] + 16}
                y1={y}
                x2={CX[1] - 16}
                y2={hy}
                stroke={on("hidden") ? "#333" : "#bbb"}
                strokeWidth="0.8"
              />
            )),
          )}
          {HID_Y.map((y, i) =>
            HID_Y.map((hy, j) => (
              <line
                key={`h-${i}-${j}`}
                x1={CX[1] + 16}
                y1={y}
                x2={CX[2] - 16}
                y2={hy}
                stroke={on("hidden") ? "#333" : "#bbb"}
                strokeWidth="0.8"
              />
            )),
          )}
          {HID_Y.map((y, i) =>
            OUT_Y.map((oy, j) => (
              <line
                key={`o-${i}-${j}`}
                x1={CX[2] + 16}
                y1={y}
                x2={CX[3] - 16}
                y2={oy}
                stroke={on("output") ? "#333" : "#bbb"}
                strokeWidth="0.8"
              />
            )),
          )}

          {INPUT_Y.map((y, i) => (
            <Neuron key={`in-${i}`} cx={CX[0]} cy={y} fill={FILL.blue} lit={on("input")} />
          ))}
          {HID_Y.map((y, i) => (
            <Neuron key={`h1-${i}`} cx={CX[1]} cy={y} fill={FILL.cyan} lit={on("hidden")} />
          ))}
          {HID_Y.map((y, i) => (
            <Neuron key={`h2-${i}`} cx={CX[2]} cy={y} fill={FILL.cyan} lit={on("hidden")} />
          ))}
          {OUT_Y.map((y, i) => (
            <Neuron key={`out-${i}`} cx={CX[3]} cy={y} fill={FILL.orange} lit={on("output")} />
          ))}

          <text x={CX[0]} y="430" textAnchor="middle" fontSize="13" fontWeight="600" fill="#222">
            Input x
          </text>
          <text x={CX[1]} y="430" textAnchor="middle" fontSize="13" fontWeight="600" fill="#222">
            Hidden
          </text>
          <text x={CX[2]} y="430" textAnchor="middle" fontSize="13" fontWeight="600" fill="#222">
            Hidden
          </text>
          <text x={CX[3]} y="430" textAnchor="middle" fontSize="13" fontWeight="600" fill="#222">
            Output y
          </text>
          <Arrow
            x1={CX[0] + 22}
            y1={455}
            x2={CX[3] - 22}
            y2={455}
            active={false}
            markerId={markerId}
          />
        </g>
      )}
    </LoopFigure>
  );
}

function Neuron({
  cx,
  cy,
  fill,
  lit,
}: {
  cx: number;
  cy: number;
  fill: string;
  lit: boolean;
}) {
  return (
    <circle
      cx={cx}
      cy={cy}
      r="16"
      fill={fill}
      stroke={lit ? "#e11d48" : "#444"}
      strokeWidth={lit ? 2.6 : 1.2}
    />
  );
}