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

const STOPS: LoopStop[] = [
  {
    x: 160,
    y: 110,
    ids: ["vision"],
    inspect: "vision",
    title: "Image patches",
    note: "The image is cut into patches. This is the new front-end; the language model still has not seen pixels.",
  },
  {
    x: 160,
    y: 170,
    ids: ["vision"],
    inspect: "vision",
    title: "ViT encoder",
    note: "A vision Transformer (often CLIP-ViT) turns patches into a sequence of embeddings.",
  },
  {
    x: 160,
    y: 270,
    ids: ["projector"],
    inspect: "projector",
    title: "Projector W",
    note: "A linear layer or small MLP maps vision width onto the LLM hidden size.",
  },
  {
    x: 500,
    y: 170,
    ids: ["decoder"],
    inspect: "decoder",
    title: "Text tokens",
    note: "The prompt is embedded as usual. Vision tokens will sit in the same residual stream.",
  },
  {
    x: 330,
    y: 360,
    ids: ["decoder", "projector"],
    inspect: "decoder",
    title: "Concat [vision ; text]",
    note: "z = LLM([W · ViT(image); text]). Cross-attention variants keep the ViT as a separate memory.",
  },
  {
    x: 330,
    y: 450,
    ids: ["decoder"],
    inspect: "decoder",
    title: "LLM decoder → logits",
    note: "The top layer is still next-token prediction. Once vision tokens are in context, the same stack can caption or follow instructions.",
  },
];

export function VlmLoop({ onSelect }: { onSelect: (id: string) => void }) {
  return (
    <LoopFigure
      stops={STOPS}
      loopMs={15000}
      viewBox="0 0 680 520"
      ariaLabel="Vision-language model architecture"
      onSelect={onSelect}
    >
      {({ on, markerId }) => (
        <g>
          <rect width="680" height="520" fill="#f7f4ec" />
          <text x="340" y="28" textAnchor="middle" fontSize="16" fontWeight="700" fill="#222">
            Vision-Language Model
          </text>
          <text x="340" y="48" textAnchor="middle" fontSize="12" fill="#555">
            z = LLM([W · ViT(image); text])
          </text>

          <text x="160" y="78" textAnchor="middle" fontSize="12" fill="#555">
            Vision
          </text>
          <g>
            {[0, 1, 2].map((r) =>
              [0, 1, 2].map((c) => (
                <rect
                  key={`${r}-${c}`}
                  x={118 + c * 28}
                  y={86 + r * 16}
                  width="24"
                  height="13"
                  rx="2"
                  fill={on("vision") ? FILL.green : "#cfe8d4"}
                  stroke="#3a7a4a"
                />
              )),
            )}
          </g>
          <Arrow x1={160} y1={140} x2={160} y2={148} active={on("vision")} markerId={markerId} />
          <FigBox
            x={70}
            y={148}
            w={180}
            h={48}
            fill={FILL.cyan}
            lit={on("vision")}
            label="ViT encoder"
            onClick={() => onSelect("vision")}
          />
          <Arrow x1={160} y1={196} x2={160} y2={242} active={on("vision") || on("projector")} markerId={markerId} />
          <FigBox
            x={70}
            y={248}
            w={180}
            h={44}
            fill={FILL.yellow}
            lit={on("projector")}
            label="Projector W"
            onClick={() => onSelect("projector")}
          />

          <text x="500" y="88" textAnchor="middle" fontSize="12" fill="#555">
            Language
          </text>
          <FigBox
            x={410}
            y={148}
            w={180}
            h={48}
            fill={FILL.pink}
            lit={on("decoder")}
            label="Text tokens"
            onClick={() => onSelect("decoder")}
          />
          <Arrow x1={500} y1={196} x2={500} y2={248} active={on("decoder")} markerId={markerId} />
          <FigBox
            x={410}
            y={248}
            w={180}
            h={44}
            fill={FILL.orange}
            lit={on("decoder")}
            label="Text embed"
            onClick={() => onSelect("decoder")}
          />

          <path
            d="M160 292 V 330 H 330"
            fill="none"
            stroke={on("projector") || on("decoder") ? "#111" : "#888"}
            strokeWidth={on("projector") ? 2.2 : 1.2}
            markerEnd={`url(#${markerId})`}
          />
          <path
            d="M500 292 V 330 H 330"
            fill="none"
            stroke={on("decoder") ? "#111" : "#888"}
            strokeWidth={on("decoder") ? 2.2 : 1.2}
            markerEnd={`url(#${markerId})`}
          />

          <FigBox
            x={220}
            y={336}
            w={220}
            h={48}
            fill={FILL.purple}
            lit={on("decoder")}
            label="[vision tokens ; text]"
            onClick={() => onSelect("decoder")}
          />
          <Arrow x1={330} y1={384} x2={330} y2={418} active={on("decoder")} markerId={markerId} />
          <FigBox
            x={200}
            y={426}
            w={260}
            h={52}
            fill={FILL.orange}
            lit={on("decoder")}
            label="LLM decoder → logits"
            onClick={() => onSelect("decoder")}
          />
        </g>
      )}
    </LoopFigure>
  );
}