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

const STOPS: LoopStop[] = [
  {
    x: 150,
    y: 90,
    ids: ["unroll"],
    inspect: "unroll",
    title: "Input x_t",
    note: "A new token (or frame) arrives. The same weights will process it as every other step.",
  },
  {
    x: 150,
    y: 200,
    ids: ["cell"],
    inspect: "cell",
    title: "Shared cell",
    note: "h_t = tanh(W_h h_{t−1} + W_x x_t). One cell, reused across time.",
  },
  {
    x: 150,
    y: 330,
    ids: ["hidden"],
    inspect: "hidden",
    title: "Hidden state h_t",
    note: "The running summary of the past. The next step reads this vector as memory.",
  },
  {
    x: 370,
    y: 200,
    ids: ["unroll", "cell"],
    inspect: "unroll",
    title: "Unroll t−1",
    note: "Training unrolls the loop into a deep net. Gradients flow back through every step (BPTT).",
  },
  {
    x: 540,
    y: 200,
    ids: ["cell", "hidden"],
    inspect: "cell",
    title: "Unroll t",
    note: "The same cell at the current time. Hidden state arrives from the left.",
  },
  {
    x: 710,
    y: 200,
    ids: ["unroll", "hidden"],
    inspect: "hidden",
    title: "Unroll t+1",
    note: "Long-range signals must survive every multiply. Then the loop returns to a new x_t.",
  },
  {
    x: 710,
    y: 400,
    ids: ["hidden"],
    inspect: "hidden",
    title: "Carry memory",
    note: "h_{t+1} is written for the next step. The rolled cell on the left is the same weights.",
  },
  {
    x: 150,
    y: 400,
    ids: ["hidden", "unroll"],
    inspect: "unroll",
    title: "Next step",
    note: "The loop returns to a new x_t with the updated hidden state.",
  },
];

export function RnnLoop({ onSelect }: { onSelect: (id: string) => void }) {
  return (
    <LoopFigure
      stops={STOPS}
      loopMs={14000}
      viewBox="0 0 860 430"
      ariaLabel="Recurrent neural network architecture, rolled and unrolled"
      onSelect={onSelect}
    >
      {({ on, markerId }) => (
        <g>
          <rect width="860" height="430" fill="#f7f4ec" />
          <text x="430" y="28" textAnchor="middle" fontSize="16" fontWeight="700" fill="#222">
            Recurrent Neural Network
          </text>
          <text x="150" y="52" textAnchor="middle" fontSize="12" fill="#555">
            Rolled
          </text>
          <text x="540" y="52" textAnchor="middle" fontSize="12" fill="#555">
            Unrolled through time
          </text>

          <FigBox
            x={80}
            y={70}
            w={140}
            h={40}
            fill={FILL.pink}
            lit={on("unroll")}
            label="x_t"
            onClick={() => onSelect("unroll")}
          />
          <Arrow x1={150} y1={110} x2={150} y2={168} active={on("unroll") || on("cell")} markerId={markerId} />
          <FigBox
            x={70}
            y={170}
            w={160}
            h={60}
            fill={FILL.cyan}
            lit={on("cell")}
            label="Shared cell"
            onClick={() => onSelect("cell")}
          />
          <path
            d="M70 200 H 40 V 120 H 150"
            fill="none"
            stroke={on("hidden") || on("cell") ? "#111" : "#888"}
            strokeWidth={on("hidden") ? 2.2 : 1.2}
            markerEnd={`url(#${markerId})`}
          />
          <text x="28" y="165" fontSize="11" fill="#333">
            h
          </text>
          <Arrow x1={150} y1={230} x2={150} y2={300} active={on("cell") || on("hidden")} markerId={markerId} />
          <FigBox
            x={80}
            y={308}
            w={140}
            h={44}
            fill={FILL.yellow}
            lit={on("hidden")}
            label="h_t"
            onClick={() => onSelect("hidden")}
          />

          <line x1="260" y1="80" x2="260" y2="380" stroke="#ccc" strokeDasharray="5 5" />

          {[
            { x: 300, t: "t−1", xid: "x_{t-1}", hid: "h_{t-1}" },
            { x: 470, t: "t", xid: "x_t", hid: "h_t" },
            { x: 640, t: "t+1", xid: "x_{t+1}", hid: "h_{t+1}" },
          ].map((col, i) => (
            <g key={col.t}>
              <text x={col.x + 70} y="78" textAnchor="middle" fontSize="12" fill="#555">
                {col.t}
              </text>
              <FigBox
                x={col.x}
                y={90}
                w={140}
                h={36}
                fill={FILL.pink}
                lit={on("unroll")}
                label={col.xid}
                onClick={() => onSelect("unroll")}
              />
              <Arrow
                x1={col.x + 70}
                y1={126}
                x2={col.x + 70}
                y2={168}
                active={on("unroll") || on("cell")}
                markerId={markerId}
              />
              <FigBox
                x={col.x}
                y={170}
                w={140}
                h={60}
                fill={FILL.cyan}
                lit={on("cell")}
                label="Cell"
                onClick={() => onSelect("cell")}
              />
              <Arrow
                x1={col.x + 70}
                y1={230}
                x2={col.x + 70}
                y2={278}
                active={on("cell") || on("hidden")}
                markerId={markerId}
              />
              <FigBox
                x={col.x}
                y={280}
                w={140}
                h={40}
                fill={FILL.yellow}
                lit={on("hidden")}
                label={col.hid}
                onClick={() => onSelect("hidden")}
              />
              {i < 2 ? (
                <Arrow
                  x1={col.x + 140}
                  y1={200}
                  x2={col.x + 168}
                  y2={200}
                  active={on("hidden") || on("cell")}
                  markerId={markerId}
                />
              ) : null}
            </g>
          ))}
        </g>
      )}
    </LoopFigure>
  );
}