js/graph.js

// graph.js — SVG renderer + pan/zoom/select for the computed Layout.

import { patternColor } from "./ir.js";

const SVG_NS = "http://www.w3.org/2000/svg";

// The HF kernels hub — a kernel (e.g. flash-attn) can be loaded to accelerate
// attention. Shown as a badge on attention nodes when flash_attention is supported.
const KERNELS_URL = "https://huggingface.co/kernels-community";

// Visual style per known edge kind (kept in sync with the sidebar legend).
const EDGE_STYLE = {
  data: { color: "#f97316", dash: null, label: "Data" },
  residual: { color: "#0ea5e9", dash: "2 5", label: "Residual" },
  mask: { color: "#8b5cf6", dash: "6 4", label: "Mask" },
  position: { color: "#10b981", dash: "6 4", label: "Position" },
  cross_attention: { color: "#ec4899", dash: "8 4", label: "Cross-Attention" },
  cache_read: { color: "#06b6d4", dash: "1 4", label: "Cache read" },
  cache_write: { color: "#f59e0b", dash: "1 4", label: "Cache write" },
};

// Colours for kinds not in the known set (e.g. cache, route, dataflow), picked
// deterministically so a given kind always renders the same colour.
const FALLBACK_PALETTE = ["#eab308", "#22d3ee", "#f472b6", "#a78bfa", "#84cc16", "#fb7185", "#60a5fa"];

function styleForKind(kind) {
  if (EDGE_STYLE[kind]) return EDGE_STYLE[kind];
  let hash = 0;
  for (const c of String(kind)) hash = (hash * 31 + c.charCodeAt(0)) >>> 0;
  const label = String(kind).replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase());
  return { color: FALLBACK_PALETTE[hash % FALLBACK_PALETTE.length], dash: "3 4", label };
}

// Sanitize an edge kind for use in DOM ids/classes.
function cssId(s) {
  return String(s).replace(/[^a-zA-Z0-9_-]/g, "_");
}

function el(tag, attrs = {}, ...children) {
  const node = document.createElementNS(SVG_NS, tag);
  for (const [k, v] of Object.entries(attrs)) {
    if (v == null) continue;
    node.setAttribute(k, v);
  }
  for (const c of children) {
    if (c == null) continue;
    node.appendChild(typeof c === "string" ? document.createTextNode(c) : c);
  }
  return node;
}

// Centre of one side of a rect ("port"). Anchoring edges at fixed ports (rather
// than at the point facing the other box's centre) keeps fan-in/out tidy: every
// child exits its bottom-centre and enters the target's top-centre.
function port(rect, side) {
  const cx = rect.x + rect.w / 2;
  const cy = rect.y + rect.h / 2;
  switch (side) {
    case "top": return { x: cx, y: rect.y };
    case "bottom": return { x: cx, y: rect.y + rect.h };
    case "left": return { x: rect.x, y: cy };
    default: return { x: rect.x + rect.w, y: cy }; // right
  }
}

// A smooth cubic path between two rects. In a top-down layout, edges that cross
// layers must route vertically (bottom→top) regardless of horizontal offset —
// so we key off vertical separation, not the dominant axis. Horizontal ports
// are used only for genuine same-row siblings (overlapping vertical extents).
function edgePath(s, t) {
  // Containment: a block feeding its own children (self_attn → q/k/v). Drop
  // straight down from the container's top edge (above the child) into the
  // child's top — reads as "the block's input fans out to these".
  const inside = (a, b) => b.x >= a.x - 1 && b.y >= a.y - 1 && b.x + b.w <= a.x + a.w + 1 && b.y + b.h <= a.y + a.h + 1;
  if (inside(s, t)) {
    const cx = t.x + t.w / 2;
    const k = Math.max(12, (t.y - s.y) * 0.4);
    return `M ${cx} ${s.y} C ${cx} ${s.y + k}, ${cx} ${t.y - k}, ${cx} ${t.y}`;
  }
  if (inside(t, s)) {
    const cx = s.x + s.w / 2;
    const k = Math.max(12, (t.y + t.h - (s.y + s.h)) * 0.4);
    return `M ${cx} ${s.y + s.h} C ${cx} ${s.y + s.h + k}, ${cx} ${t.y + t.h - k}, ${cx} ${t.y + t.h}`;
  }
  const gap = 4;
  let p1, p2, axis;
  if (t.y >= s.y + s.h - gap) {
    axis = "v";
    p1 = port(s, "bottom");
    p2 = port(t, "top"); // target is below → downward flow
  } else if (s.y >= t.y + t.h - gap) {
    axis = "v";
    p1 = port(s, "top");
    p2 = port(t, "bottom"); // target is above → upward flow
  } else {
    axis = "h";
    const toRight = t.x + t.w / 2 >= s.x + s.w / 2;
    p1 = port(s, toRight ? "right" : "left");
    p2 = port(t, toRight ? "left" : "right");
  }
  if (axis === "v") {
    const k = Math.max(18, Math.abs(p2.y - p1.y) * 0.4);
    const sgn = Math.sign(p2.y - p1.y) || 1;
    return `M ${p1.x} ${p1.y} C ${p1.x} ${p1.y + sgn * k}, ${p2.x} ${p2.y - sgn * k}, ${p2.x} ${p2.y}`;
  }
  const k = Math.max(18, Math.abs(p2.x - p1.x) * 0.4);
  const sgn = Math.sign(p2.x - p1.x) || 1;
  return `M ${p1.x} ${p1.y} C ${p1.x + sgn * k} ${p1.y}, ${p2.x - sgn * k} ${p2.y}, ${p2.x} ${p2.y}`;
}

class Graph {
  constructor(svg, callbacks = {}) {
    this.svg = svg;
    this.cb = callbacks;
    this.tx = 0;
    this.ty = 0;
    this.scale = 1;
    this.selectedId = null;
    this.layout = null;

    this.overrides = new Map(); // id -> {dx,dy} manual position nudges
    this.disp = new Map(); // id -> displayed {x,y,w,h} (base + overrides)
    this._raf = null;

    this._defs();
    this.viewport = el("g", { class: "viewport" });
    this.svg.appendChild(this.viewport);

    this._bindPointer();
  }

  // Clear all manual node positions.
  clearOverrides() {
    this.overrides.clear();
  }
  resetPositions() {
    this.clearOverrides();
    if (this.layout) this.render(this.layout, this.opts);
  }

  _defs() {
    this.defs = el("defs");
    this._markers = new Set();
    this.svg.appendChild(this.defs);
    Object.keys(EDGE_STYLE).forEach((k) => this._ensureMarker(k));
    // Shared circular clip for org avatars (object-bounding-box → any size).
    const clip = el(
      "clipPath",
      { id: "avatar-clip", clipPathUnits: "objectBoundingBox" },
      el("circle", { cx: "0.5", cy: "0.5", r: "0.5" })
    );
    this.defs.appendChild(clip);
  }

  // Lazily create an arrowhead marker for any edge kind (incl. unknown ones).
  _ensureMarker(kind) {
    const id = `arrow-${cssId(kind)}`;
    if (this._markers.has(id)) return id;
    this._markers.add(id);
    const marker = el(
      "marker",
      {
        id,
        viewBox: "0 0 10 10",
        refX: "9",
        refY: "5",
        markerWidth: "7",
        markerHeight: "7",
        orient: "auto-start-reverse",
      },
      el("path", { d: "M0,0 L10,5 L0,10 z", fill: styleForKind(kind).color })
    );
    this.defs.appendChild(marker);
    return id;
  }

  // --- Pan / zoom -----------------------------------------------------------

  _apply() {
    this.viewport.setAttribute(
      "transform",
      `translate(${this.tx} ${this.ty}) scale(${this.scale})`
    );
    if (!this._suppressView && this.cb.onView) {
      this.cb.onView({ tx: this.tx, ty: this.ty, scale: this.scale });
    }
  }

  // Mirror another graph's view (for synced side-by-side panels).
  setView(v) {
    this._suppressView = true;
    this.tx = v.tx;
    this.ty = v.ty;
    this.scale = v.scale;
    this._apply();
    this._suppressView = false;
  }

  _scheduleRender() {
    if (this._raf) return;
    const raf =
      typeof requestAnimationFrame !== "undefined" ? requestAnimationFrame : (f) => f();
    this._raf = raf(() => {
      this._raf = null;
      if (this.layout) this.render(this.layout, this.opts);
    });
  }

  _bindPointer() {
    // Either panning the canvas or dragging a single node/subtree.
    let mode = null; // 'pan' | 'node'
    let dragId = null;
    let sx = 0;
    let sy = 0;
    let stx = 0;
    let sty = 0;
    let baseOverride = { dx: 0, dy: 0 };
    let moved = false;

    this.svg.addEventListener("pointerdown", (e) => {
      if (e.target.closest("[data-toggle]")) return; // toggle handles itself
      moved = false;
      sx = e.clientX;
      sy = e.clientY;
      const nodeEl = e.target.closest("[data-node]");
      if (nodeEl) {
        mode = "node";
        dragId = nodeEl.getAttribute("data-node");
        const o = this.overrides.get(dragId) || { dx: 0, dy: 0 };
        baseOverride = { dx: o.dx, dy: o.dy };
      } else {
        mode = "pan";
        stx = this.tx;
        sty = this.ty;
      }
      this.svg.classList.add("grabbing");
      this.svg.setPointerCapture(e.pointerId);
    });

    this.svg.addEventListener("pointermove", (e) => {
      if (!mode) return;
      const ddx = e.clientX - sx;
      const ddy = e.clientY - sy;
      if (Math.abs(ddx) + Math.abs(ddy) > 3) moved = true;
      if (mode === "pan") {
        this.tx = stx + ddx;
        this.ty = sty + ddy;
        this._apply();
      } else if (mode === "node" && moved) {
        // Drag distance is in screen px; convert to graph units.
        this.overrides.set(dragId, {
          dx: baseOverride.dx + ddx / this.scale,
          dy: baseOverride.dy + ddy / this.scale,
        });
        this._scheduleRender();
      }
    });

    const end = (e) => {
      if (mode === "node" && !moved && dragId != null) {
        this.cb.onSelect && this.cb.onSelect(dragId);
      }
      mode = null;
      dragId = null;
      this.svg.classList.remove("grabbing");
      try {
        this.svg.releasePointerCapture(e.pointerId);
      } catch (_) {}
    };
    this.svg.addEventListener("pointerup", end);
    this.svg.addEventListener("pointercancel", end);

    this.svg.addEventListener(
      "wheel",
      (e) => {
        e.preventDefault();
        const rect = this.svg.getBoundingClientRect();
        const mx = e.clientX - rect.left;
        const my = e.clientY - rect.top;
        const factor = Math.exp(-e.deltaY * 0.0015);
        const next = Math.max(0.2, Math.min(2.5, this.scale * factor));
        // Zoom toward cursor.
        this.tx = mx - (mx - this.tx) * (next / this.scale);
        this.ty = my - (my - this.ty) * (next / this.scale);
        this.scale = next;
        this._apply();
      },
      { passive: false }
    );
  }

  fit() {
    if (!this.layout) return;
    const rect = this.svg.getBoundingClientRect();
    // Don't compute a transform while the SVG is hidden (0×0) — it would shrink
    // the content to nothing and leave a blank view once shown again.
    if (rect.width < 1 || rect.height < 1) return;
    const margin = 48;
    const sw = Math.max(1, rect.width - margin * 2);
    const sh = Math.max(1, rect.height - margin * 2);
    const s = Math.min(sw / this.layout.width, sh / this.layout.height, 1.4);
    this.scale = Number.isFinite(s) && s > 0 ? s : 1;
    this.tx = (rect.width - this.layout.width * this.scale) / 2;
    this.ty = margin;
    this._apply();
  }

  // --- Rendering ------------------------------------------------------------

  // Accumulate manual overrides down the containment chain so dragging a
  // container moves its whole subtree with it.
  _computeDisp() {
    this.disp = new Map();
    const off = (id) => {
      const r = this.layout.rectById.get(id);
      const o = this.overrides.get(id) || { dx: 0, dy: 0 };
      const p = r && r.parentId ? off(r.parentId) : { dx: 0, dy: 0 };
      return { dx: o.dx + p.dx, dy: o.dy + p.dy };
    };
    for (const r of this.layout.placed) {
      const o = off(r.id);
      this.disp.set(r.id, { x: r.x + o.dx, y: r.y + o.dy, w: r.w, h: r.h });
    }
  }

  // A copy of a placed rect with displayed (override-adjusted) coordinates.
  _dr(r) {
    const d = this.disp.get(r.id);
    return { ...r, x: d.x, y: d.y };
  }

  render(layout, opts) {
    this.layout = layout;
    this.opts = opts;
    this._computeDisp();
    this.viewport.replaceChildren();

    const frames = el("g", { class: "layer-frames" });
    const edgesG = el("g", { class: "layer-edges" });
    const nodesG = el("g", { class: "layer-nodes" });

    for (const r of layout.placed) {
      if (r.isRoot) continue;
      if (r.kind === "container") frames.appendChild(this._container(this._dr(r)));
    }
    for (const e of layout.visibleEdges(opts.edgeKinds)) {
      edgesG.appendChild(this._edge(e));
    }
    for (const r of layout.placed) {
      if (r.isRoot || r.kind === "container") continue;
      nodesG.appendChild(this._leaf(this._dr(r)));
    }

    // Edges last → on top of nodes, so a wire is never hidden behind an
    // intervening block (you can trace it to its real target). The layer is
    // non-interactive so nodes underneath stay clickable/draggable.
    this.viewport.appendChild(frames);
    this.viewport.appendChild(nodesG);
    this.viewport.appendChild(edgesG);
    this._applySelection();
  }

  _shared(id) {
    return this.opts.sharedIds && this.opts.sharedIds.has(id) ? " shared" : "";
  }

  // Merged-diff origin class ("origin-a" / "origin-b"; "both" gets no class).
  _origin(id) {
    if (!this.opts.origin) return "";
    const o = this.opts.origin.get(id);
    return o === "a" || o === "b" ? ` origin-${o}` : "";
  }

  _container(r) {
    const cls = (r.isRepeat ? "group group-repeat" : "group group-plain") + this._shared(r.id) + this._origin(r.id);
    const g = el("g", { class: cls, "data-node": r.id });
    if (r.isRepeat) {
      const deck = this._deck(r, 12); // keep the depth impression when expanded
      if (deck) g.appendChild(deck);
    }
    g.appendChild(
      el("rect", {
        class: "group-box",
        x: r.x,
        y: r.y,
        width: r.w,
        height: r.h,
        rx: 12,
      })
    );
    // Header bar.
    g.appendChild(
      el("rect", {
        class: "group-header",
        x: r.x,
        y: r.y,
        width: r.w,
        height: r.headerH,
        rx: 12,
      })
    );
    const label = this.ir().label(r.id, this.opts.fields);
    const sub = r.isRepeat ? null : this.ir().kindLabel(r.node && r.node.kind);
    // With a schedule strip in the header, pin the title near the top; else centre.
    const titleY = r.hasSchedule ? r.y + 17 : r.y + r.headerH / 2 + 4;
    g.appendChild(
      el(
        "text",
        { class: "group-title", x: r.x + 12, y: titleY },
        r.isRepeat ? label : label + (sub ? `  ·  ${sub}` : "")
      )
    );
    if (r.hasSchedule) {
      const sched = this.ir().scheduleForRepeat(r.node, this.opts.fields);
      if (sched) g.appendChild(this._scheduleStrip(sched, r.x, r.y + r.headerH - 12, r.w));
    }
    const ck = r.node && r.node.kind;
    if (ck === "attention" || ck === "cross_attention") {
      const tags = this._backendTags(r.x + r.w, r.y + r.headerH / 2 - 8);
      if (tags) g.appendChild(tags);
    }
    // Kernel badge for a kernelizable container (e.g. a MoE MLP).
    if (this.opts.showKernels !== false) {
      const kb = this._kernelBadge(r.node, r.x + r.w, r.y + r.headerH / 2 - 8);
      if (kb) g.appendChild(kb);
    }
    if (r.isRepeat) g.appendChild(this._toggle(r, true));
    return g;
  }

  _leaf(r) {
    const kind = r.node ? r.node.kind : "input";
    if (r.kind === "input") {
      // Pseudo endpoint; class by prefix so state:* looks distinct from input:*.
      const prefix = String(r.id).includes(":") ? String(r.id).split(":")[0] : "input";
      const g = el("g", { class: `node node-pseudo node-pseudo-${cssId(prefix)}`, "data-node": r.id });
      g.appendChild(
        el("rect", { x: r.x, y: r.y, width: r.w, height: r.h, rx: r.h / 2 })
      );
      g.appendChild(
        el(
          "text",
          { class: "node-title", x: r.x + r.w / 2, y: r.y + r.h / 2 + 4, "text-anchor": "middle" },
          this.ir().label(r.id, this.opts.fields)
        )
      );
      return g;
    }

    const g = el("g", {
      class: `node ${r.isRepeat ? "node-repeat" : "node-leaf"}${this._shared(r.id)}${this._origin(r.id)}`,
      "data-node": r.id,
      "data-kind": kind,
    });
    if (r.isRepeat) {
      const deck = this._deck(r, 10); // stacked "deck" — depth = layers, coloured by schedule
      if (deck) g.appendChild(deck);
    }
    g.appendChild(el("rect", { class: "node-box", x: r.x, y: r.y, width: r.w, height: r.h, rx: 10 }));
    // Left accent stripe, inset past the box's corner radius so it stays inside.
    g.appendChild(el("rect", { class: "kind-bar", x: r.x, y: r.y + 10, width: 4, height: Math.max(6, r.h - 20), rx: 2 }));
    g.appendChild(el("circle", { class: "kind-dot", cx: r.x + 16, cy: r.y + 19, r: 4 }));

    const title = this.ir().label(r.id, this.opts.fields);
    g.appendChild(el("text", { class: "node-title", x: r.x + 28, y: r.y + 23 }, title));
    const sub = this.ir().kindLabel(kind);
    g.appendChild(el("text", { class: "node-sub", x: r.x + 28, y: r.y + 41 }, sub));
    // Shape / attribute caption (from the IR) — the extra density.
    if (r.info) {
      g.appendChild(el("text", { class: "node-info", x: r.x + 28, y: r.y + 58 }, r.info));
    }

    if (r.hasSchedule) {
      const sched = this.ir().scheduleForRepeat(r.node, this.opts.fields);
      if (sched) g.appendChild(this._scheduleStrip(sched, r.x, r.y + r.h - 13, r.w));
    }
    if (kind === "attention" || kind === "cross_attention") {
      const tags = this._backendTags(r.x + r.w, r.y + 8);
      if (tags) g.appendChild(tags);
    }
    // Projection shard-direction glyph (top-right).
    if (this.opts.showTP !== false && kind === "projection" && r.node.attributes && r.node.attributes.tp) {
      g.appendChild(this._tpGlyph(r.node.attributes.tp, r.x + r.w - 20, r.y + 9));
    }
    // Kernel badge for any kernelizable leaf (RMSNorm, …).
    if (this.opts.showKernels !== false) {
      const kb = this._kernelBadge(r.node, r.x + r.w, r.y + 8);
      if (kb) g.appendChild(kb);
    }
    // "tied" badge on the word embedding when weights are shared with the LM head.
    if (this.opts.tied && this.opts.ir && this.opts.ir.isWordEmbedding(r.node)) {
      g.appendChild(this._tiedBadge(r.x + r.w, r.y + 8));
    }
    if (r.isRepeat) g.appendChild(this._toggle(r, false));
    return g;
  }

  // Right-aligned attention-backend tags for an attention node: the faster
  // implementations that could replace this layer, plus a Kernels badge when
  // Flash Attention is available (a kernel can be loaded to integrate it).
  _backendTags(rightEdge, y) {
    const caps = this.opts.ir && this.opts.ir.capabilities;
    const backends = caps && caps.attention_backends;
    if (!Array.isArray(backends) || !backends.length) return null;
    const g = el("g", { class: "attn-tags" });
    let cx = rightEdge - 10;
    const tag = (label, cls, href) => {
      const w = Math.round(12 + label.length * 6.0);
      cx -= w;
      const t = el("g", { class: `attn-tag ${cls}` });
      t.appendChild(el("rect", { x: cx, y, width: w, height: 16, rx: 8 }));
      t.appendChild(el("text", { class: "attn-tag-txt", x: cx + w / 2, y: y + 11.5, "text-anchor": "middle" }, label));
      if (href) {
        t.addEventListener("click", (ev) => {
          ev.stopPropagation();
          if (typeof window !== "undefined" && window.open) window.open(href, "_blank", "noopener");
        });
      }
      g.appendChild(t);
      cx -= 5;
    };
    ["flex_attention", "flash_attention", "sdpa", "eager"]
      .filter((b) => backends.includes(b))
      .forEach((b) =>
        tag(
          b.startsWith("flash") ? "⚡ flash" : b.replace("_attention", ""),
          b.startsWith("flash") ? "flash" : b === "eager" ? "muted" : "std"
        )
      );
    return g;
  }

  // Unified ⚡ Kernels badge for any node the generator marked as kernelizable
  // (attributes.kernel). Links to the kernel's Hub repo. Right-aligned at rightEdge.
  _kernelBadge(node, rightEdge, y) {
    const k = this.opts.ir && this.opts.ir.nodeKernel(node);
    if (!k) return null;
    const open = (href) => (ev) => {
      ev.stopPropagation();
      if (typeof window !== "undefined" && window.open) window.open(href, "_blank", "noopener");
    };
    const g = el("g", {});
    const avatars = this.opts.avatars;
    const orgs = [...new Set(k.repos.map((r) => String(r).split("/")[0]))];
    // Org avatars (deduped), right-to-left.
    let ax = rightEdge - 6;
    for (const org of orgs) {
      const url = avatars && avatars.get(org);
      if (!url) continue;
      ax -= 18;
      const a = el("g", { class: "org-avatar" });
      a.appendChild(el("title", {}, `${k.name} · ${org}`));
      a.appendChild(el("image", { href: url, x: ax, y, width: 18, height: 18, "clip-path": "url(#avatar-clip)", preserveAspectRatio: "xMidYMid slice" }));
      a.appendChild(el("circle", { class: "org-avatar-ring", cx: ax + 9, cy: y + 9, r: 9 }));
      a.addEventListener("click", open(`https://huggingface.co/${org}`));
      g.appendChild(a);
      ax -= 4;
    }
    // Lightning bolt, always shown (left of the avatars).
    ax -= 12;
    const b = el("g", { class: "kernel-bolt-badge" });
    b.appendChild(el("title", {}, `${k.name}${k.repos.length ? " · " + k.repos.join(", ") : ""}`));
    b.appendChild(el("text", { x: ax + 6, y: y + 13, "text-anchor": "middle" }, "⚡"));
    b.addEventListener("click", open(k.repos.length ? `https://huggingface.co/${k.repos[0]}` : KERNELS_URL));
    g.appendChild(b);
    return g;
  }

  // Tensor-parallel shard-direction glyph: vertical bars = column-parallel
  // (output split), horizontal bars = row-parallel (input split).
  _tpGlyph(tp, x, y) {
    const g = el("g", { class: `tp-glyph tp-${tp}` });
    g.appendChild(
      el("title", {}, tp === "colwise" ? "column-parallel · splits output" : "row-parallel · splits input")
    );
    const size = 12;
    const bar = 2.5;
    const step = 4.5;
    for (let i = 0; i < 3; i++) {
      g.appendChild(
        tp === "colwise"
          ? el("rect", { x: x + i * step, y, width: bar, height: size, rx: 1 })
          : el("rect", { x, y: y + i * step, width: size, height: bar, rx: 1 })
      );
    }
    return g;
  }

  // "🔗 tied" badge — the input embedding shares weights with the output head.
  _tiedBadge(rightEdge, y) {
    const label = "🔗 tied";
    const w = Math.round(14 + label.length * 6.0);
    const x = rightEdge - w - 8;
    const g = el("g", { class: "tied-badge" });
    g.appendChild(el("title", {}, "tied word embeddings — input embedding weights shared with the output (LM head)"));
    g.appendChild(el("rect", { x, y, width: w, height: 16, rx: 8 }));
    g.appendChild(el("text", { class: "attn-tag-txt", x: x + w / 2, y: y + 11.5, "text-anchor": "middle" }, label));
    return g;
  }

  // A stacked "deck" behind a repeat block conveying its depth: one card per
  // layer, offset down-right, coloured by the attention schedule when present
  // (so sliding×5+full reads as 5 amber slivers then 1 blue) or neutral. Drawn
  // back-to-front so nearer layers sit on top; the block box covers the rest.
  _deck(r, rx) {
    const geo = this.opts.ir && this.opts.ir.deckGeometry(r.node, this.opts.fields);
    if (!geo || !geo.cards.length) return null;
    const g = el("g", { class: "deck" });
    // Back-to-front: deepest band first, so nearer (earlier) layers sit on top.
    for (let i = geo.cards.length - 1; i >= 0; i--) {
      const c = geo.cards[i];
      g.appendChild(
        el("rect", {
          // Scheduled cards colour via inline fill; neutral ones via the class
          // (a CSS `fill` rule would otherwise override the inline attribute).
          class: c.pattern ? "deck-card" : "deck-card neutral",
          x: r.x + c.off * 0.6,
          y: r.y + c.off,
          width: r.w,
          height: r.h,
          rx,
          fill: c.pattern ? patternColor(c.pattern) : undefined,
        })
      );
    }
    return g;
  }

  // A row of per-layer attention-schedule cells. When there's room it prefixes
  // an "attn/layer" label and appends a legend (swatch + pattern name) so the
  // colours are self-explanatory; otherwise it just centres the bare cells.
  _scheduleStrip(sched, x, y, w) {
    const g = el("g", { class: "sched-strip" });
    const n = sched.length;
    const short = (p) => String(p).replace(/_attention$/, "");
    const distinct = [...new Set(sched)];

    const cw = 8;
    const gap = 1;
    const cellsW = n * cw + (n - 1) * gap;
    const label = "attn/layer";
    const labelW = label.length * 5.6 + 8;
    const legendItemW = (p) => 11 + 4 + short(p).length * 6 + 12;
    const legendW = distinct.reduce((s, p) => s + legendItemW(p), 0);
    const full = labelW + cellsW + 10 + legendW;

    const drawCells = (cx0) => {
      let cx = cx0;
      for (let i = 0; i < n; i++) {
        const c = el("rect", { x: cx, y, width: cw - gap, height: 8, rx: 1.5, fill: patternColor(sched[i]) });
        c.appendChild(el("title", {}, `layer ${i}: ${short(sched[i])}`));
        g.appendChild(c);
        cx += cw + gap;
      }
      return cx;
    };

    if (full + 16 <= w) {
      // label · cells · legend, centred as a unit
      let cur = x + (w - full) / 2;
      g.appendChild(el("text", { class: "sched-strip-label", x: cur, y: y + 7.5 }, label));
      cur += labelW;
      cur = drawCells(cur) + 8;
      for (const p of distinct) {
        g.appendChild(el("rect", { x: cur, y: y - 1, width: 11, height: 10, rx: 2, fill: patternColor(p) }));
        g.appendChild(el("text", { class: "sched-strip-label", x: cur + 15, y: y + 7.5 }, short(p)));
        cur += legendItemW(p);
      }
    } else {
      // narrow (collapsed block): bare centred cells with per-cell tooltips
      const fitCw = Math.max(2, Math.min(9, (w - 20) / n));
      const fg = fitCw > 4 ? 1 : 0;
      const tw = n * fitCw + (n - 1) * fg;
      let cx = x + (w - tw) / 2;
      for (let i = 0; i < n; i++) {
        const c = el("rect", { x: cx, y, width: Math.max(1, fitCw - fg), height: 8, rx: 1.5, fill: patternColor(sched[i]) });
        c.appendChild(el("title", {}, `layer ${i}: ${short(sched[i])}`));
        g.appendChild(c);
        cx += fitCw + fg;
      }
    }
    return g;
  }

  _toggle(r, expanded) {
    const size = 18;
    const cx = r.x + r.w - size / 2 - 8;
    const cy = r.y + (expanded ? r.headerH / 2 : 16);
    const g = el("g", { class: "toggle", "data-toggle": r.id });
    g.appendChild(el("circle", { cx, cy, r: size / 2 }));
    g.appendChild(
      el("path", {
        class: "toggle-icon",
        d: expanded
          ? `M ${cx - 4} ${cy} H ${cx + 4}`
          : `M ${cx - 4} ${cy} H ${cx + 4} M ${cx} ${cy - 4} V ${cy + 4}`,
      })
    );
    g.addEventListener("click", (ev) => {
      ev.stopPropagation();
      this.cb.onToggle && this.cb.onToggle(r.id);
    });
    return g;
  }

  _edge(e) {
    const s = this.disp.get(e.source);
    const t = this.disp.get(e.target);
    const d = edgePath(s, t);
    const style = styleForKind(e.kind);
    const markerId = this._ensureMarker(e.kind);
    const observed = e.raw && e.raw.observed_forward;
    // Merged-diff: colour edges that exist in only one model.
    let originCls = "";
    let stroke = style.color;
    if (this.opts.edgeOrigin && e.raw) {
      const o = this.opts.edgeOrigin.get(`${e.raw.source}|${e.raw.target}|${e.raw.kind}`);
      if (o === "a" || o === "b") {
        originCls = ` origin-${o}`;
        stroke = o === "a" ? "#38bdf8" : "#f87171";
      }
    }
    const g = el("g", {
      class: `edge edge-${cssId(e.kind)}${observed ? " edge-observed" : ""}${originCls}`,
      "data-src": e.source,
      "data-dst": e.target,
    });
    // Background halo (reads as "passing over" a block) + the coloured line.
    g.appendChild(el("path", { class: "edge-halo", d, fill: "none" }));
    g.appendChild(
      el("path", {
        class: "edge-line",
        d,
        fill: "none",
        stroke,
        "stroke-dasharray": style.dash,
        "marker-end": `url(#${markerId})`,
      })
    );
    return g;
  }

  // --- Selection ------------------------------------------------------------

  setSelected(id) {
    this.selectedId = id;
    this._applySelection();
  }

  _applySelection() {
    this.viewport.querySelectorAll("[data-node]").forEach((n) => {
      n.classList.toggle("selected", n.getAttribute("data-node") === this.selectedId);
    });
    this.viewport.querySelectorAll(".edge").forEach((ed) => {
      const on =
        this.selectedId &&
        (ed.getAttribute("data-src") === this.selectedId ||
          ed.getAttribute("data-dst") === this.selectedId);
      ed.classList.toggle("edge-hi", !!on);
      ed.classList.toggle("edge-dim", !!this.selectedId && !on);
    });
  }

  ir() {
    return this.opts.ir;
  }
}

export { Graph, EDGE_STYLE, styleForKind };