import { useMemo } from "react";

import { formatNumber } from "./formula";
import { edgePath, layoutGraph, OUTPUT, SLOT_OFFSET, type PlacedEdge, type PlacedShape } from "./layout";
import MathText from "./MathText";
import StepFlow, { TOKEN_BASELINE } from "./StepFlow";
import type { Flow, Timeline } from "./Timeline";
import type { Graph, State } from "./types";

// Small graphs keep their size rather than filling the plot
const MIN_VIEW_WIDTH = 560;
const MIN_VIEW_HEIGHT = 240;
// The baseline of an edge's name, relative to the edge
const ABOVE_EDGE = -5;

function SlotText({ x, y, number }: { x: number; y: number; number: number }) {
  return (
    <text x={x} y={y + TOKEN_BASELINE} textAnchor="middle" strokeWidth={3} paintOrder="stroke" className="fill-ink stroke-white font-mono text-xs">
      {formatNumber(number)}
    </text>
  );
}

// A shape with its name or operator in the middle, its value above and its derivative below
function Shape({ shape, current, value, derivative }: { shape: PlacedShape; current: boolean; value: number | undefined; derivative: number | undefined }) {
  const fill = current ? "fill-marker-soft" : shape.kind === "op" ? "fill-white" : "fill-paper";
  return (
    <g data-node={shape.id}>
      {shape.kind === "op" ? (
        <ellipse cx={shape.x} cy={shape.y} rx={shape.width / 2} ry={shape.height / 2} strokeWidth={1.5} className={`stroke-ink ${fill}`} />
      ) : (
        <rect
          x={shape.x - shape.width / 2}
          y={shape.y - shape.height / 2}
          width={shape.width}
          height={shape.height}
          strokeWidth={1.5}
          className={`stroke-ink ${fill}`}
        />
      )}
      <text x={shape.x} y={shape.y + 5} textAnchor="middle" className="fill-ink font-mono text-sm font-medium">
        <MathText text={shape.label} svg />
      </text>
      {value !== undefined && <SlotText x={shape.x} y={shape.y - SLOT_OFFSET} number={value} />}
      {derivative !== undefined && <SlotText x={shape.x} y={shape.y + SLOT_OFFSET} number={derivative} />}
    </g>
  );
}

export default function ComputationGraph({
  graph,
  state,
  flow,
  timeline,
}: {
  graph: Graph;
  state: State | undefined;
  flow: Flow | null;
  timeline: Timeline;
}) {
  const layout = useMemo(() => layoutGraph(graph), [graph]);
  const viewWidth = Math.max(layout.width, MIN_VIEW_WIDTH);
  const viewHeight = Math.max(layout.height, MIN_VIEW_HEIGHT);
  const step = state?.step;
  const backward = step?.phase === "backward";
  // The edges into the current node, or the edge from f when the backward pass starts
  const isActive = (edge: PlacedEdge) =>
    step !== undefined &&
    (backward && step.partials.length === 0 ? edge.to === OUTPUT : edge.to === step.node && edge.to !== OUTPUT);
  // Drawn last, so they sit on top
  const edges = [...layout.edges.filter((edge) => !isActive(edge)), ...layout.edges.filter(isActive)];

  return (
    <svg
      className="min-h-[240px] w-full flex-1"
      viewBox={`${(layout.width - viewWidth) / 2} ${(layout.height - viewHeight) / 2} ${viewWidth} ${viewHeight}`}
      role="img"
      aria-label="Computational graph of the function"
    >
      <defs>
        <marker id="arrowhead" viewBox="0 0 8 8" refX="7" refY="4" markerWidth="8" markerHeight="8" orient="auto">
          <path d="M 0 0 L 8 4 L 0 8 z" className="fill-ink-muted" />
        </marker>
        <marker id="arrowhead-active" viewBox="0 0 8 8" refX="7" refY="4" markerWidth="6" markerHeight="6" orient="auto-start-reverse">
          <path d="M 0 0 L 8 4 L 0 8 z" className="fill-ink" />
        </marker>
      </defs>
      {edges.map((edge) => {
        const active = isActive(edge);
        // The backward pass flows against the edge, so its arrow points back
        const flipped = active && backward;
        return (
          <path
            key={edge.id}
            d={edgePath(edge)}
            fill="none"
            strokeWidth={active ? 3 : 1.5}
            markerStart={flipped ? "url(#arrowhead-active)" : undefined}
            markerEnd={flipped ? undefined : active ? "url(#arrowhead-active)" : "url(#arrowhead)"}
            className={active ? "stroke-ink" : "stroke-ink-muted"}
          />
        );
      })}
      {layout.shapes.map((shape) => (
        <Shape
          key={shape.id}
          shape={shape}
          current={shape.id === step?.node}
          // f shows what the output node holds
          value={state?.values[shape.id === OUTPUT ? graph.output : shape.id]}
          derivative={state?.derivatives[shape.id === OUTPUT ? graph.output : shape.id]}
        />
      ))}
      {layout.edges.map((edge) => (
        <text
          key={edge.id}
          data-edge={edge.from}
          x={edge.labelAt.x}
          y={edge.labelAt.y + ABOVE_EDGE}
          // On a slanted edge, the name moves to the side the edge leaves clear
          textAnchor={edge.labelSlope === 0 ? "middle" : edge.labelSlope < 0 ? "end" : "start"}
          strokeWidth={4}
          paintOrder="stroke"
          className="fill-ink stroke-white font-mono text-xs"
        >
          <MathText text={edge.from} svg />
        </text>
      ))}
      {flow && <StepFlow key={flow.id} graph={graph} layout={layout} flow={flow} timeline={timeline} />}
    </svg>
  );
}