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

const H = [
  { y: 110, label: "h₁" },
  { y: 190, label: "h₂" },
  { y: 270, label: "h₃" },
  { y: 350, label: "h₄" },
];

const STOPS: LoopStop[] = [
  {
    x: 120,
    y: 190,
    ids: ["score"],
    inspect: "score",
    title: "Encoder states",
    note: "Every source position keeps its hidden state. Attention will read from all of them.",
  },
  {
    x: 520,
    y: 90,
    ids: ["score"],
    inspect: "score",
    title: "Decoder state s_t",
    note: "The query: how relevant is each encoder state to the token being generated now?",
  },
  {
    x: 330,
    y: 200,
    ids: ["score"],
    inspect: "score",
    title: "Score(s_t, h_i)",
    note: "A small net (or a dot product) measures alignment between the decoder and each source position.",
  },
  {
    x: 330,
    y: 300,
    ids: ["weights"],
    inspect: "weights",
    title: "Softmax → α_i",
    note: "Scores become a distribution over the source. That heatmap is the attention.",
  },
  {
    x: 520,
    y: 250,
    ids: ["context"],
    inspect: "context",
    title: "Context c_t",
    note: "c_t = Σ α_i h_i. A dynamic memory read — no single bottleneck vector.",
  },
  {
    x: 520,
    y: 360,
    ids: ["context"],
    inspect: "context",
    title: "Next token",
    note: "The decoder uses c_t (and s_t) to emit the next symbol, then the loop advances t.",
  },
  {
    x: 520,
    y: 410,
    ids: ["context"],
    inspect: "context",
    title: "Advance t",
    note: "The decoder steps forward and queries the same encoder memory again.",
  },
  {
    x: 120,
    y: 410,
    ids: ["score"],
    inspect: "score",
    title: "Back to source",
    note: "Encoder states stay; only the decoder time index changes.",
  },
];

export function AttentionLoop({ onSelect }: { onSelect: (id: string) => void }) {
  return (
    <LoopFigure
      stops={STOPS}
      loopMs={14000}
      viewBox="0 0 700 430"
      ariaLabel="Bahdanau encoder-decoder attention architecture"
      onSelect={onSelect}
    >
      {({ on, markerId }) => (
        <g>
          <rect width="700" height="430" fill="#f7f4ec" />
          <text x="350" y="28" textAnchor="middle" fontSize="16" fontWeight="700" fill="#222">
            Attention (Bahdanau 2014)
          </text>
          <text x="350" y="48" textAnchor="middle" fontSize="12" fill="#555">
            c_t = Σ softmax(score(s_t, h_i)) h_i
          </text>

          <text x="120" y="78" textAnchor="middle" fontSize="12" fill="#555">
            Encoder
          </text>
          {H.map((h) => (
            <g key={h.label}>
              <FigBox
                x={50}
                y={h.y - 20}
                w={140}
                h={40}
                fill={FILL.blue}
                lit={on("score")}
                label={h.label}
                onClick={() => onSelect("score")}
              />
              <path
                d={`M190 ${h.y} H 250 V 200 H 258`}
                fill="none"
                stroke={on("score") ? "#111" : "#999"}
                strokeWidth={on("score") ? 1.8 : 1}
              />
            </g>
          ))}

          <FigBox
            x={260}
            y={170}
            w={140}
            h={60}
            fill={FILL.orange}
            lit={on("score")}
            label="Score"
            onClick={() => onSelect("score")}
          />
          <Arrow x1={330} y1={230} x2={330} y2={268} active={on("score") || on("weights")} markerId={markerId} />
          <FigBox
            x={260}
            y={276}
            w={140}
            h={48}
            fill={FILL.green}
            lit={on("weights")}
            label="Softmax α"
            onClick={() => onSelect("weights")}
          />

          <Arrow x1={400} y1={90} x2={330} y2={170} active={on("score")} markerId={markerId} />
          <FigBox
            x={450}
            y={70}
            w={160}
            h={40}
            fill={FILL.pink}
            lit={on("score")}
            label="Decoder s_t"
            onClick={() => onSelect("score")}
          />
          <Arrow
            x1={400}
            y1={300}
            x2={450}
            y2={250}
            active={on("weights") || on("context")}
            markerId={markerId}
          />
          <FigBox
            x={450}
            y={228}
            w={160}
            h={44}
            fill={FILL.yellow}
            lit={on("context")}
            label="Context c_t"
            onClick={() => onSelect("context")}
          />
          <Arrow x1={530} y1={272} x2={530} y2={328} active={on("context")} markerId={markerId} />
          <FigBox
            x={450}
            y={336}
            w={160}
            h={44}
            fill={FILL.purple}
            lit={on("context")}
            label="Next token"
            onClick={() => onSelect("context")}
          />

          <text x="120" y="410" textAnchor="middle" fontSize="11" fill="#666">
            Source positions
          </text>
          <text x="530" y="410" textAnchor="middle" fontSize="11" fill="#666">
            Target step t
          </text>
        </g>
      )}
    </LoopFigure>
  );
}