Blame for frontends/react/src/OptimisationPlot 2.tsx

import { useRef, useState, type RefObject } from "react";
import { usePlotlyClick } from "@elvis/ui";
import Plot from "react-plotly.js";
import type { Camera, Data, Layout } from "plotly.js";

import type { Trajectory } from "./useOptimisation";
import type { FunctionValues, Range, SurfaceValues } from "./types";

interface OptimisationPlotProps {
  values: FunctionValues | null;
  trajectories: Trajectory[];
  xRange: RefObject<Range>;
  yRange: RefObject<Range>;
  onViewChange: (x: Range, y: Range) => void;
  // Called with data coordinates when the plot is clicked; null ignores clicks
  onClick: ((x: number, y: number) => void) | null;
  // With two variables: a 3D surface instead of a contour map
  surface: boolean;
  arrows: boolean;
}

const CONFIG = { scrollZoom: true, displayModeBar: true, doubleClick: false, responsive: true } as const;
const LEGEND = { x: 0.02, y: 0.98, xanchor: "left", yanchor: "top", bgcolor: "rgba(255,255,255,0.8)" } as const;
const INK = "#1B2A4A";

export default function OptimisationPlot({ values, trajectories, xRange, yRange, onViewChange, onClick, surface, arrows }: OptimisationPlotProps) {
  const plotRef = useRef<Plot | null>(null);
  // The 3D camera as the student last left it. Plotly resets it on each redraw
  // while stepping (uirevision alone doesn't keep it), so it's passed back in.
  const [camera, setCamera] = useState<Partial<Camera> | undefined>(undefined);
  // While the student drags the surface, the frame they started on; a redraw
  // from stepping would otherwise cancel the drag
  const [held, setHeld] = useState<{ data: Data[]; layout: Partial<Layout> } | null>(null);
  // el exists at runtime only
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  usePlotlyClick(() => (plotRef.current as any)?.el, surface ? null : onClick);

  const twoVariables = values !== null && "y" in values;
  const showLegend = trajectories.length > 1;

  if (twoVariables && surface) {
    const data = [surfaceTrace(values), ...trajectories.flatMap((t) => pathTraces(t, true))];
    const layout: Partial<Layout> = {
      // Keep the camera while stepping
      uirevision: "surface",
      showlegend: showLegend,
      legend: LEGEND,
      margin: { t: 20, r: 20, b: 20, l: 20 },
      scene: {
        camera,
        aspectmode: "cube",
        xaxis: { title: { text: "x" }, range: [...xRange.current] },
        yaxis: { title: { text: "y" }, range: [...yRange.current] },
        zaxis: { title: { text: "f(x, y)" } },
      },
    };
    return (
      <div
        className="flex-1 min-h-[320px]"
        onPointerDown={() => {
          setHeld({ data, layout });
          // Catch up once the drag ends, after Plotly has reported the new camera
          window.addEventListener("pointerup", () => setTimeout(() => setHeld(null), 50), { once: true });
        }}
      >
        <Plot
          key="surface"
          data={held?.data ?? data}
          layout={held?.layout ?? layout}
          onRelayout={(event) => {
            const moved = (event as { "scene.camera"?: Partial<Camera> })["scene.camera"];
            if (moved) {
              setCamera(moved);
            }
          }}
          config={CONFIG}
          style={{ width: "100%", height: "100%" }}
        />
      </div>
    );
  }

  const data: Data[] = [];
  if (values && "y" in values) {
    data.push(contourTrace(values));
    if (arrows) {
      data.push(...arrowTraces(values, xRange.current, yRange.current));
    }
  } else if (values) {
    data.push({
      x: values.x,
      y: values.f,
      type: "scatter",
      mode: "lines",
      line: { color: INK, width: 2 },
      name: "f(x)",
      showlegend: false,
      hovertemplate: "x = %{x:.4g}<br>f(x) = %{y:.4g}<extra></extra>",
    });
  }
  data.push(...trajectories.flatMap((t) => pathTraces(t, false, twoVariables)));

  const layout: Partial<Layout> = {
    dragmode: "pan",
    // Keep the user's pan and zoom while stepping; a new range (e.g. from a preset) still applies
    uirevision: "view",
    showlegend: showLegend,
    legend: LEGEND,
    margin: { t: 40, r: 40, b: 40, l: 40 },
    xaxis: { title: { text: "x" }, range: [...xRange.current] },
    yaxis: { title: { text: twoVariables ? "y" : "f(x)" }, range: [...yRange.current] },
  };

  return (
    <div className="flex-1 min-h-[320px]">
      <Plot
        key="flat"
        ref={plotRef}
        data={data}
        layout={layout}
        onRelayout={(event) => {
          // eslint-disable-next-line @typescript-eslint/no-explicit-any
          const full = (plotRef.current as any)?.el?._fullLayout;
          const x = (full?.xaxis?.range ?? [event["xaxis.range[0]"], event["xaxis.range[1]"]]) as Range;
          const y = (full?.yaxis?.range ?? [event["yaxis.range[0]"], event["yaxis.range[1]"]]) as Range;
          const changed = x[0] !== xRange.current[0] || x[1] !== xRange.current[1] || y[0] !== yRange.current[0] || y[1] !== yRange.current[1];
          if (changed && x.every(Number.isFinite) && y.every(Number.isFinite)) {
            onViewChange([x[0], x[1]], [y[0], y[1]]);
          }
        }}
        config={CONFIG}
        style={{ width: "100%", height: "100%", minWidth: 0 }}
      />
    </div>
  );
}

function contourTrace(values: SurfaceValues): Data {
  return {
    x: values.x,
    y: values.y,
    z: values.f,
    type: "contour",
    colorscale: "Viridis",
    contours: { coloring: "heatmap", showlines: false },
    colorbar: { thickness: 12, title: { text: "f" } },
    hovertemplate: "x = %{x:.3g}, y = %{y:.3g}<br>f = %{z:.4g}<extra></extra>",
  } as Data;
}

function surfaceTrace(values: SurfaceValues): Data {
  return {
    x: values.x,
    y: values.y,
    z: values.f,
    type: "surface",
    colorscale: "Viridis",
    opacity: 0.85,
    colorbar: { thickness: 12, title: { text: "f" } },
    hovertemplate: "x = %{x:.3g}, y = %{y:.3g}<br>f = %{z:.4g}<extra></extra>",
  } as Data;
}

// A run's path and, larger, its current point
function pathTraces(t: Trajectory, threeD: boolean, twoVariables = true): Data[] {
  const xs = t.path.map((p) => p.position[0]);
  const ys = twoVariables ? t.path.map((p) => p.position[1]) : t.path.map((p) => p.value);
  const fs = t.path.map((p) => p.value);
  const hovertemplate = twoVariables
    ? "x = %{x:.4g}, y = %{y:.4g}<br>f = %{customdata:.4g}<extra>%{fullData.name}</extra>"
    : "x = %{x:.4g}<br>f(x) = %{y:.4g}<extra>%{fullData.name}</extra>";
  const last = t.path.length - 1;
  if (threeD) {
    return [
      { x: xs, y: ys, z: fs, type: "scatter3d", mode: "lines+markers", name: t.label, legendgroup: t.label,
        line: { color: t.color, width: 5 }, marker: { color: t.color, size: 2.5 }, hoverinfo: "skip" },
      { x: xs.slice(last), y: ys.slice(last), z: fs.slice(last), type: "scatter3d", mode: "markers", name: t.label,
        legendgroup: t.label, showlegend: false, marker: { color: t.color, size: 6, line: { color: "white", width: 1 } }, hoverinfo: "skip" },
    ] as Data[];
  }
  return [
    { x: xs, y: ys, customdata: fs, type: "scatter", mode: "lines+markers", name: t.label, legendgroup: t.label,
      line: { color: t.color, width: 2 }, marker: { color: t.color, size: 7 }, hovertemplate },
    { x: xs.slice(last), y: ys.slice(last), customdata: fs.slice(last), type: "scatter", mode: "markers", name: t.label,
      legendgroup: t.label, showlegend: false, marker: { color: t.color, size: 14, line: { color: "white", width: 2 } }, hovertemplate },
  ] as Data[];
}

// The gradient as arrows (pointing uphill), scaled so a typical arrow spans
// most of the gap between grid points
function arrowTraces({ arrows }: SurfaceValues, [x0, x1]: Range, [y0, y1]: Range): Data[] {
  // Arrows per side of the grid
  const side = Math.sqrt(arrows.x.length);
  // Lengths as fractions of the view, so x and y scale alike
  const lengths = arrows.u.map((u, i) => Math.hypot((u ?? NaN) / (x1 - x0), (arrows.v[i] ?? NaN) / (y1 - y0)));
  const finite = lengths.filter((l) => Number.isFinite(l) && l > 0).sort((a, b) => a - b);
  const typical = finite[Math.floor(0.9 * (finite.length - 1))] ?? 1;

  const shaftX: (number | null)[] = [];
  const shaftY: (number | null)[] = [];
  const headX: number[] = [];
  const headY: number[] = [];
  const angles: number[] = [];
  const hover: number[][] = [];
  arrows.x.forEach((x, i) => {
    const u = arrows.u[i];
    const v = arrows.v[i];
    if (u === null || v === null || !(lengths[i] > 0)) {
      return;
    }
    // At most 0.9 of a grid cell (1 / side of the view)
    const scale = (0.9 * Math.min(1, lengths[i] / typical)) / (lengths[i] * side);
    const dx = u * scale;
    const dy = v * scale;
    const y = arrows.y[i];
    shaftX.push(x, x + dx, null);
    shaftY.push(y, y + dy, null);
    headX.push(x + dx);
    headY.push(y + dy);
    // Plotly turns markers clockwise from up, in screen space
    angles.push((Math.atan2(dx / (x1 - x0), dy / (y1 - y0)) * 180) / Math.PI);
    hover.push([x, y, u, v]);
  });

  const color = "rgba(255, 255, 255, 0.85)";
  return [
    { x: shaftX, y: shaftY, type: "scatter", mode: "lines", line: { color, width: 1 }, hoverinfo: "skip", showlegend: false },
    {
      x: headX, y: headY, customdata: hover, type: "scatter", mode: "markers", showlegend: false,
      marker: { symbol: "triangle-up", size: 6, color, angle: angles, angleref: "up" },
      hovertemplate: "x = %{customdata[0]:.3g}, y = %{customdata[1]:.3g}<br>∇f = (%{customdata[2]:.3g}, %{customdata[3]:.3g})<extra></extra>",
    },
  ] as Data[];
}