web/app/ProgramGraph.tsx

"use client";

import { useMemo } from "react";
import ReactFlow, {
  Edge,
  Handle,
  MarkerType,
  Node,
  NodeProps,
  Position,
  ReactFlowProvider,
} from "reactflow";
import "reactflow/dist/style.css";
import "./ProgramGraph.css";
import { ReprNode, parseProgramRepr } from "../lib/programRepr";

// Palette tokens (must match tailwind.config.ts).
const C = {
  bg: "#FAFAF7",
  card: "#FCFBF8",
  ink: "#23303A",
  muted: "#6E7F8C",
  border: "#ECEAE4",
  accent: "#3A6B7E",
  highlight: "#BC6B2E",
  matrix: "#F1EFEA",
  score: "#FBF1E6",
};

export type GenePill = { id: string; symbol?: string | null; matched?: boolean };

interface ProgramGraphProps {
  programRepr: string;
  outputLabel?: string;
  /** opaque_id -> { symbol, matched } once revealed. */
  reveal?: Record<string, { symbol: string; matched: boolean }>;
}

// ---------- custom node components ----------------------------------------

function DataPillNode({ data }: NodeProps<{ label: string }>) {
  return (
    <div
      className="pg-node"
      style={{
        background: C.matrix,
        border: `1px solid ${C.border}`,
        color: C.ink,
        borderRadius: 999,
        padding: "8px 14px",
        fontSize: 12,
        whiteSpace: "nowrap",
      }}
    >
      <Handle type="source" position={Position.Right} className="pg-handle" />
      <Handle type="target" position={Position.Left} className="pg-handle" />
      {data.label}
    </div>
  );
}

function ScorePillNode({ data }: NodeProps<{ label: string }>) {
  return (
    <div
      className="pg-node"
      style={{
        background: C.score,
        border: `1px solid ${C.highlight}`,
        color: C.highlight,
        borderRadius: 999,
        padding: "8px 14px",
        fontSize: 12,
        fontWeight: 500,
        whiteSpace: "nowrap",
      }}
    >
      <Handle type="source" position={Position.Right} className="pg-handle" />
      <Handle type="target" position={Position.Left} className="pg-handle" />
      {data.label}
    </div>
  );
}

function VerbNode({
  data,
}: NodeProps<{ label: string; sub?: string; pills?: GenePill[] }>) {
  return (
    <div
      className="pg-node"
      style={{
        background: C.card,
        border: `1px solid ${C.accent}`,
        color: C.ink,
        borderRadius: 8,
        padding: "10px 14px",
        minWidth: 130,
        boxShadow: "0 1px 0 rgba(35,48,58,0.02)",
      }}
    >
      <Handle type="source" position={Position.Right} className="pg-handle" />
      <Handle type="target" position={Position.Left} className="pg-handle" />
      <div style={{ fontSize: 12, fontWeight: 600, color: C.ink }}>
        {data.label}
      </div>
      {data.sub && (
        <div style={{ marginTop: 2, fontSize: 10, color: C.muted }}>
          {data.sub}
        </div>
      )}
      {data.pills && data.pills.length > 0 && (
        <div
          style={{
            marginTop: 6,
            display: "flex",
            flexWrap: "wrap",
            gap: 4,
            maxWidth: 220,
          }}
        >
          {data.pills.map((p) => {
            const display = p.symbol ?? p.id;
            const matched = p.matched;
            return (
              <span
                key={p.id}
                title={p.symbol ? `${p.id} → ${p.symbol}` : p.id}
                style={{
                  fontFamily:
                    'ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace',
                  fontSize: 10,
                  padding: "2px 6px",
                  borderRadius: 4,
                  background: matched ? C.highlight : C.bg,
                  color: matched ? C.card : C.ink,
                  border: `1px solid ${matched ? C.highlight : C.border}`,
                  fontWeight: matched ? 600 : 400,
                  lineHeight: 1.3,
                }}
              >
                {display}
              </span>
            );
          })}
        </div>
      )}
    </div>
  );
}

function OutputNode({ data }: NodeProps<{ label: string }>) {
  return (
    <div
      className="pg-node"
      style={{
        background: C.highlight,
        color: "#FFFFFF",
        border: `1px solid ${C.highlight}`,
        borderRadius: 999,
        padding: "10px 16px",
        fontSize: 12,
        fontWeight: 600,
        whiteSpace: "nowrap",
        boxShadow: "0 1px 0 rgba(35,48,58,0.04)",
      }}
    >
      <Handle type="target" position={Position.Left} className="pg-handle" />
      {data.label}
    </div>
  );
}

function TierGroupNode({
  data,
}: NodeProps<{ label: string; variant: "tier1" | "tier2"; width: number; height: number }>) {
  const isTier2 = data.variant === "tier2";
  return (
    <div
      style={{
        width: data.width,
        height: data.height,
        borderRadius: 12,
        background: isTier2 ? "transparent" : "#F1F6F7",
        border: isTier2 ? "1px dashed #BC6B2E" : "1px solid #3A6B7E",
        position: "relative",
        pointerEvents: "none",
      }}
    >
      <span
        style={{
          position: "absolute",
          top: 6,
          left: 10,
          padding: "0 6px",
          fontSize: 10,
          letterSpacing: 0.4,
          fontWeight: 600,
          color: isTier2 ? "#BC6B2E" : "#3A6B7E",
          background: isTier2 ? C.bg : "#F1F6F7",
        }}
      >
        {data.label}
      </span>
    </div>
  );
}

const nodeTypes = {
  dataPill: DataPillNode,
  scorePill: ScorePillNode,
  verb: VerbNode,
  output: OutputNode,
  tierGroup: TierGroupNode,
};

// ---------- layout --------------------------------------------------------

const COL_W = 200;     // horizontal spacing between layers
const ROW_H = 95;      // vertical spacing between sibling lanes

interface Built {
  nodes: Node[];
  edges: Edge[];
  width: number;
  height: number;
}

interface Counter {
  id: number;
}

function nid(c: Counter, prefix: string): string {
  c.id += 1;
  return `${prefix}${c.id}`;
}

function pillsFor(
  features: string[],
  reveal?: Record<string, { symbol: string; matched: boolean }>,
): GenePill[] {
  return features.map((id) => {
    const r = reveal?.[id];
    return r ? { id, symbol: r.symbol, matched: r.matched } : { id };
  });
}

/** Recursively build nodes/edges for a subtree rooted at ``tree``.
 * Returns the SubLayout — its own width / height bounding box and the
 * (x,y) of its right-output anchor (where the parent connects). */
type SubLayout = {
  rootId: string;
  width: number;   // span in x
  height: number;  // span in y
  anchorY: number; // vertical centre relative to the layout
};

function buildSubtree(
  tree: ReprNode,
  reveal: ProgramGraphProps["reveal"],
  c: Counter,
  out: { nodes: Node[]; edges: Edge[] },
  xOffset: number,
  yOffset: number,
): SubLayout {
  if (tree.kind === "M") {
    const id = nid(c, "M");
    out.nodes.push({
      id,
      type: "dataPill",
      position: { x: xOffset, y: yOffset },
      data: { label: "Expression matrix" },
      draggable: false,
    });
    return { rootId: id, width: COL_W, height: ROW_H, anchorY: yOffset + 18 };
  }

  if (tree.kind === "Select") {
    const child = buildSubtree(tree.matrix, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "Sel");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: {
        label: "Select",
        sub: `${tree.features.length} gene${tree.features.length === 1 ? "" : "s"}`,
        pills: pillsFor(tree.features, reveal),
      },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    const width = child.width + COL_W;
    return {
      rootId: id,
      width,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "Reduce") {
    const child = buildSubtree(tree.matrix, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "Red");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: { label: "Reduce", sub: tree.agg },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    return {
      rootId: id,
      width: child.width + COL_W,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "Combine") {
    const left = buildSubtree(tree.left, reveal, c, out, xOffset, yOffset);
    const right = buildSubtree(
      tree.right,
      reveal,
      c,
      out,
      xOffset,
      yOffset + left.height + 20,
    );
    const id = nid(c, "Cmb");
    const maxChildWidth = Math.max(left.width, right.width);
    const x = xOffset + maxChildWidth;
    const centre = (left.anchorY + right.anchorY) / 2;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y: centre - 18 },
      data: { label: "Combine", sub: tree.op },
      draggable: false,
    });
    out.edges.push(edge(`${left.rootId}->${id}`, left.rootId, id));
    out.edges.push(edge(`${right.rootId}->${id}`, right.rootId, id));
    return {
      rootId: id,
      width: maxChildWidth + COL_W,
      height: left.height + right.height + 20,
      anchorY: centre,
    };
  }

  if (tree.kind === "Search") {
    const child = buildSubtree(tree.matrix, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "Srch");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: { label: "Search", sub: `top-${tree.k} (bounded)` },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    return {
      rootId: id,
      width: child.width + COL_W,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "Split") {
    const child = buildSubtree(tree.inner, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "Splt");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: { label: "Split", sub: tree.predicate },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    return {
      rootId: id,
      width: child.width + COL_W,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "FitApply") {
    const child = buildSubtree(tree.inner, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "FitA");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: { label: "Fit → Apply", sub: `target: ${tree.target}` },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    return {
      rootId: id,
      width: child.width + COL_W,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "Associate") {
    const child = buildSubtree(tree.inner, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "Assoc");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: {
        label: "Associate",
        sub: `${tree.assocKind} · target: ${tree.target}`,
      },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    return {
      rootId: id,
      width: child.width + COL_W,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "Effect") {
    const child = buildSubtree(tree.inner, reveal, c, out, xOffset, yOffset);
    const id = nid(c, "Eff");
    const x = xOffset + child.width;
    const y = child.anchorY - 18;
    out.nodes.push({
      id,
      type: "verb",
      position: { x, y },
      data: {
        label: "Effect",
        sub: `${tree.kind} · adjust: stage, age`,
      },
      draggable: false,
    });
    out.edges.push(edge(`${child.rootId}->${id}`, child.rootId, id));
    return {
      rootId: id,
      width: child.width + COL_W,
      height: Math.max(child.height, ROW_H),
      anchorY: child.anchorY,
    };
  }

  if (tree.kind === "Fit") {
    // Legacy v1 fallback: render each child subtree, then a "Fit" verb,
    // then the output node.
    const stacks: SubLayout[] = [];
    let yCursor = yOffset;
    for (const child of tree.children) {
      const lay = buildSubtree(child, reveal, c, out, xOffset, yCursor);
      stacks.push(lay);
      yCursor += lay.height + 20;
    }
    const childMaxWidth = Math.max(...stacks.map((s) => s.width));
    const fitId = nid(c, "Fit");
    const centre =
      stacks.reduce((a, s) => a + s.anchorY, 0) /
      Math.max(stacks.length, 1);
    out.nodes.push({
      id: fitId,
      type: "verb",
      position: { x: xOffset + childMaxWidth, y: centre - 18 },
      data: { label: "Fit", sub: tree.output || "target" },
      draggable: false,
    });
    for (const s of stacks) {
      out.edges.push(edge(`${s.rootId}->${fitId}`, s.rootId, fitId));
    }
    return {
      rootId: fitId,
      width: childMaxWidth + COL_W,
      height: stacks.reduce((a, s) => a + s.height + 20, 0),
      anchorY: centre,
    };
  }

  // Unknown — render as a muted box.
  const id = nid(c, "U");
  out.nodes.push({
    id,
    type: "verb",
    position: { x: xOffset, y: yOffset },
    data: { label: "?", sub: tree.text.slice(0, 32) },
    draggable: false,
  });
  return { rootId: id, width: COL_W, height: ROW_H, anchorY: yOffset + 18 };
}

function edge(id: string, source: string, target: string): Edge {
  return {
    id,
    source,
    target,
    type: "smoothstep",
    style: { stroke: "#B9B6AE", strokeWidth: 1.4 },
    markerEnd: {
      type: MarkerType.ArrowClosed,
      color: "#B9B6AE",
      width: 14,
      height: 14,
    },
  };
}

function buildGraph(props: ProgramGraphProps): Built {
  const out = { nodes: [] as Node[], edges: [] as Edge[] };
  const c: Counter = { id: 0 };
  let tree: ReprNode;
  try {
    tree = parseProgramRepr(props.programRepr);
  } catch (e) {
    out.nodes.push({
      id: "err",
      type: "verb",
      position: { x: 20, y: 20 },
      data: { label: "couldn't parse program", sub: String(e) },
      draggable: false,
    });
    return { ...out, width: 320, height: 120 };
  }

  if (tree.kind === "Combine") {
    return buildCombineRoot(tree, props, c, out);
  }

  const lay = buildSubtree(tree, props.reveal, c, out, 20, 20);
  // Append a terminal output pill for typed (engine_v2) trees that end in
  // a Vector but have no explicit output node. The legacy Fit root path
  // already adds a "Fit" verb labelled with the output, so we skip there.
  if (tree.kind !== "Fit") {
    const id = nid(c, "Out");
    const x = lay.width + 40;
    const y = lay.anchorY - 14;
    out.nodes.push({
      id,
      type: "output",
      position: { x, y },
      data: { label: props.outputLabel ?? "score" },
      draggable: false,
    });
    out.edges.push(edge(`${lay.rootId}->${id}`, lay.rootId, id));
    return { ...out, width: lay.width + COL_W + 40, height: lay.height + 60 };
  }
  return { ...out, width: lay.width + 40, height: lay.height + 60 };
}

// ---------- Combine root: Tier-1 boxes per side + Tier-2 dashed wrap -----

const SUBTREE_X = 40;        // left padding inside Tier-2 wrapper
const SUBTREE_Y0 = 44;       // top padding (room for Tier-2 label)
const TIER_GAP = 56;         // vertical gap between Tier-1 groups
const TIER_PAD = 18;         // padding around node bboxes inside Tier-1
const TIER_LABEL_H = 22;     // visual room for "Tier-1" label

function buildCombineRoot(
  tree: ReprNode & { kind: "Combine" },
  props: ProgramGraphProps,
  c: Counter,
  out: { nodes: Node[]; edges: Edge[] },
): Built {
  const leftStart = out.nodes.length;
  const left = buildSubtree(tree.left, props.reveal, c, out, SUBTREE_X, SUBTREE_Y0);
  const leftRange: [number, number] = [leftStart, out.nodes.length];

  const rightStart = out.nodes.length;
  const rightY = SUBTREE_Y0 + left.height + TIER_GAP;
  const right = buildSubtree(tree.right, props.reveal, c, out, SUBTREE_X, rightY);
  const rightRange: [number, number] = [rightStart, out.nodes.length];

  const leftBbox = bboxFromNodes(out.nodes, leftRange);
  const rightBbox = bboxFromNodes(out.nodes, rightRange);

  // Tier-1 boxes are positioned to contain each side's bbox.
  const leftTier = boxForBbox(leftBbox, "Tier-1", "tier1");
  const rightTier = boxForBbox(rightBbox, "Tier-1", "tier1");

  // Place Combine + Output to the right of the wider subtree.
  const subtreeMaxRight = Math.max(leftTier.maxX, rightTier.maxX);
  const cmbX = subtreeMaxRight + 30;
  const centreY = (left.anchorY + right.anchorY) / 2;
  const cmbId = nid(c, "Cmb");
  out.nodes.push({
    id: cmbId,
    type: "verb",
    position: { x: cmbX, y: centreY - 18 },
    data: { label: "Combine", sub: tree.op },
    draggable: false,
  });
  out.edges.push(edge(`${left.rootId}->${cmbId}`, left.rootId, cmbId));
  out.edges.push(edge(`${right.rootId}->${cmbId}`, right.rootId, cmbId));

  const outX = cmbX + COL_W;
  const outId = nid(c, "Out");
  out.nodes.push({
    id: outId,
    type: "output",
    position: { x: outX, y: centreY - 14 },
    data: { label: props.outputLabel ?? "score" },
    draggable: false,
  });
  out.edges.push(edge(`${cmbId}->${outId}`, cmbId, outId));

  // Tier-2 dashed wrap around EVERYTHING (left + right + Combine + Output).
  const allBbox = unionBboxes([
    leftTier,
    rightTier,
    { minX: cmbX, minY: centreY - 22, maxX: outX + 170, maxY: centreY + 22 },
  ]);
  const tier2 = boxForBbox(allBbox, "Tier-2 program", "tier2");

  // Insert tier-group nodes at the START so they render BEHIND everything.
  out.nodes.unshift(tierGroupNode(c, leftTier, "Tier-1", "tier1"));
  out.nodes.unshift(tierGroupNode(c, rightTier, "Tier-1", "tier1"));
  out.nodes.unshift(tierGroupNode(c, tier2, "Tier-2 program", "tier2"));

  const totalWidth = tier2.maxX + 24;
  const totalHeight = tier2.maxY + 24;
  return { ...out, width: totalWidth, height: totalHeight };
}

interface Bbox {
  minX: number;
  minY: number;
  maxX: number;
  maxY: number;
}

function bboxFromNodes(nodes: Node[], [start, end]: [number, number]): Bbox {
  let minX = Infinity;
  let minY = Infinity;
  let maxX = -Infinity;
  let maxY = -Infinity;
  for (let i = start; i < end; i++) {
    const n = nodes[i];
    const { x, y } = n.position;
    const { w, h } = approxSize(n);
    if (x < minX) minX = x;
    if (y < minY) minY = y;
    if (x + w > maxX) maxX = x + w;
    if (y + h > maxY) maxY = y + h;
  }
  return { minX, minY, maxX, maxY };
}

function approxSize(n: Node): { w: number; h: number } {
  // Conservative bounds for each node type. Verb nodes with many gene
  // pills get taller — estimate from the data payload.
  const t = n.type;
  if (t === "dataPill" || t === "scorePill" || t === "output") {
    return { w: 170, h: 38 };
  }
  if (t === "verb") {
    const pills: unknown = (n.data as { pills?: unknown })?.pills;
    const pillCount = Array.isArray(pills) ? pills.length : 0;
    const pillRows = pillCount > 0 ? Math.ceil(pillCount / 4) : 0;
    return { w: 220, h: 56 + pillRows * 22 };
  }
  return { w: 170, h: 56 };
}

function unionBboxes(list: Bbox[]): Bbox {
  return {
    minX: Math.min(...list.map((b) => b.minX)),
    minY: Math.min(...list.map((b) => b.minY)),
    maxX: Math.max(...list.map((b) => b.maxX)),
    maxY: Math.max(...list.map((b) => b.maxY)),
  };
}

function boxForBbox(b: Bbox, _label: string, _variant: "tier1" | "tier2"): Bbox {
  // Same shape; bbox-with-padding used both for the visual tier rect and
  // for downstream layout.
  return {
    minX: b.minX - TIER_PAD,
    minY: b.minY - TIER_LABEL_H - 6,
    maxX: b.maxX + TIER_PAD,
    maxY: b.maxY + TIER_PAD,
  };
}

function tierGroupNode(
  c: Counter,
  b: Bbox,
  label: string,
  variant: "tier1" | "tier2",
): Node {
  return {
    id: nid(c, "tier-"),
    type: "tierGroup",
    position: { x: b.minX, y: b.minY },
    data: {
      label,
      variant,
      width: b.maxX - b.minX,
      height: b.maxY - b.minY,
    },
    draggable: false,
    selectable: false,
    style: { zIndex: 0 },
  };
}

export default function ProgramGraph(props: ProgramGraphProps) {
  const { nodes, edges, width, height } = useMemo(
    () => buildGraph(props),
    [props],
  );
  const containerHeight = Math.max(220, height);
  // Remount-key so ReactFlow's one-shot `fitView` runs again every
  // time the laid-out program changes (winner ↔ candidate, pasted,
  // dataset swap). Without this a larger Tier-2 graph stays at the
  // previous zoom and clips on the right/bottom.
  const fitKey = `${width}x${height}:${nodes.length}:${edges.length}`;

  return (
    <div
      style={{ height: containerHeight, background: "transparent", borderRadius: 8 }}
      className="pg-canvas"
    >
      <ReactFlowProvider>
        <ReactFlow
          key={fitKey}
          nodes={nodes}
          edges={edges}
          nodeTypes={nodeTypes}
          fitView
          fitViewOptions={{ padding: 0.18 }}
          nodesDraggable={false}
          nodesConnectable={false}
          edgesFocusable={false}
          panOnDrag={false}
          panOnScroll={false}
          zoomOnScroll={false}
          zoomOnPinch={false}
          zoomOnDoubleClick={false}
          preventScrolling={false}
          proOptions={{ hideAttribution: true }}
          // A floor of 0.4 was too strict for full 2-tier programs;
          // 0.2 lets fitView always scale the whole tree to fit the
          // card width without horizontal overflow.
          minZoom={0.2}
          maxZoom={1.2}
        />
      </ReactFlowProvider>
    </div>
  );
}