Blame for frontends/react/src/useOptimisation 2.ts

import { useCallback, useEffect, useRef, useState } from "react";
import { usePythonStatus } from "@elvis/pyodide";
import { useStepper, type ParamValue } from "@elvis/ui";

import { algorithmById } from "./algorithms";
import { RUN_COLORS } from "./colors";
import { ALL_PRESETS, detectVariables, PRESETS, type FunctionPreset } from "./functions";
import { python } from "./python";
import type { FunctionValues, Point, Range, RunSettings, State, Variables } from "./types";

// After typing stops, wait this long before rerunning
const TYPING_DELAY_MS = 500;
const STEP_INTERVAL_MS = 200;
// Runs shown at once (the current one plus comparisons), one per colour
export const MAX_RUNS = RUN_COLORS.length;

// A run Python is stepping; id 1 is the main run
export interface Run {
  id: number;
  settings: RunSettings;
  color: string;
  label: string;
}

// A run's points up to the stepper's iteration
export interface Trajectory extends Run {
  path: Point[];
  // Why it ended, once it has
  stopped: string | null;
}

// The range covering the finite values, with 10% padding on each side
function paddedRange(values: (number | null)[]): Range | null {
  const finite = values.filter((v): v is number => v !== null && Number.isFinite(v));
  if (finite.length === 0) {
    return null;
  }
  const min = Math.min(...finite);
  const max = Math.max(...finite);
  const pad = max > min ? 0.1 * (max - min) : 1;
  return [min - pad, max + pad];
}

// A coordinate picked on the plot, to about 3 significant figures of the view
function roundTo([a, b]: Range, value: number): string {
  const decimals = Math.max(0, 2 - Math.floor(Math.log10(Math.abs(b - a))));
  return String(Number(value.toFixed(decimals)));
}

// One optimiser's settings as the student edits them, with the start kept as typed
interface RunDraft {
  algorithm: string;
  params: Record<string, ParamValue>;
  startText: string[];
}

// A run added in the Compare tab: its own settings, a fixed number and a colour
export interface Comparison extends RunDraft {
  id: number;
  color: string;
}

function newDraft(start: number[]): RunDraft {
  return { algorithm: "gradient_descent", params: {}, startText: start.map(String) };
}

// Edits one optimiser's draft through update; onTyping runs before typed changes
function makeEditor(draft: RunDraft, update: (change: (d: RunDraft) => RunDraft) => void, onTyping: () => void) {
  const start = draft.startText.map((t) => (t.trim() === "" ? NaN : Number(t)));
  const settings: RunSettings = { algorithm: draft.algorithm, start, params: draft.params };
  return {
    algorithm: draft.algorithm,
    params: draft.params,
    startText: draft.startText,
    startValid: start.every(Number.isFinite),
    settings,
    setStartText(startText: string[]) {
      update((d) => ({ ...d, startText }));
    },
    changeAlgorithm(algorithm: string) {
      update((d) => ({ ...d, algorithm, params: {} }));
    },
    changeParams(params: Record<string, ParamValue>) {
      onTyping();
      update((d) => ({ ...d, params }));
    },
    changeStart(index: number, text: string) {
      onTyping();
      update((d) => ({ ...d, startText: d.startText.map((t, i) => (i === index ? text : t)) }));
    },
  };
}

export type RunEditor = ReturnType<typeof makeEditor>;

export default function useOptimisation() {
  const pythonStatus = usePythonStatus(python);
  const first = PRESETS[1][0];

  const [expression, setExpression] = useState(first.expression);
  // 1 or 2, from whether the function uses y
  const variables = detectVariables(expression);
  const [mainDraft, setMainDraft] = useState(() => newDraft(first.start));
  // Runs added in the Compare tab, stepped alongside the main one
  const [compared, setCompared] = useState<Comparison[]>([]);
  // The next comparison's number (the main run is 1)
  const nextId = useRef(2);
  const [surface, setSurface] = useState(false);
  const [showArrows, setShowArrows] = useState(false);

  const [functionError, setFunctionError] = useState<string | null>(null);
  const [runError, setRunError] = useState<string | null>(null);
  const [values, setValues] = useState<FunctionValues | null>(null);
  const [runs, setRuns] = useState<Run[]>([]);

  // The plot's view; changed by panning and zooming, or by picking a preset
  const xRange = useRef<Range>(first.x);
  const yRange = useRef<Range>([-1, 1]);
  // With one variable, fit the y axis to the next curve
  const fitY = useRef(true);
  // Delay before the next rerun: TYPING_DELAY_MS after a keystroke, else 0
  const delay = useRef(0);
  const typing = () => {
    delay.current = TYPING_DELAY_MS;
  };
  // Every run reruns from iteration 0 as any of them is edited
  const main = makeEditor(mainDraft, setMainDraft, typing);
  const comparisonEditors = compared.map((c) =>
    makeEditor(c, (change) => setCompared((cs) => cs.map((x) => (x.id === c.id ? { ...x, ...change(x) } : x))), typing),
  );
  // The function Python has, as "variables\nexpression"
  const appliedFunction = useRef<string | null>(null);
  // Bumped by each rerun, so older responses are ignored
  const version = useRef(0);
  // Every state fetched since the last rerun, by iteration
  const states = useRef<State[]>([]);

  const allRuns: Run[] = [
    { id: 1, settings: main.settings, color: RUN_COLORS[0].value, label: `1 · ${algorithmById(main.algorithm).label}` },
    ...compared.map((c, i) => ({
      id: c.id,
      settings: comparisonEditors[i].settings,
      color: c.color,
      label: `${c.id} · ${algorithmById(c.algorithm).label}`,
    })),
  ];
  const allSettings = allRuns.map((run) => run.settings);
  const allValid = main.startValid && comparisonEditors.every((e) => e.startValid);

  const fetchState = useCallback(async (iteration: number) => {
    const v = version.current;
    const result = await python.call("handle_state", { iteration });
    if (!result.ok) {
      throw new Error(result.error.message);
    }
    if (result.data && v === version.current) {
      states.current[iteration] = result.data;
    }
    return result.data;
  }, []);
  const step = useStepper(fetchState, { intervalMs: STEP_INTERVAL_MS });
  const { stepper } = step;

  async function refreshValues() {
    const result = await python.call("handle_function_values", { x_range: xRange.current, y_range: yRange.current });
    if (!result.ok) {
      console.error(`Python ${result.error.kind} error: ${result.error.message}`);
      return;
    }
    const data = result.data;
    if (data && !("y" in data) && fitY.current) {
      yRange.current = paddedRange(data.f) ?? yRange.current;
    }
    fitY.current = false;
    setValues(data);
  }

  // Send the function (if it changed) and the runs to Python, then step from 0
  async function apply(functionKey: string, next: Run[]) {
    const v = ++version.current;
    stepper.pause();
    if (functionKey !== appliedFunction.current) {
      appliedFunction.current = null;
      const result = await python.call("handle_set_function", { variables, expression });
      if (v !== version.current) {
        return;
      }
      setFunctionError(result.ok ? null : result.error.message);
      if (!result.ok) {
        setValues(null);
        setRuns([]);
        states.current = [];
        await stepper.restart();
        return;
      }
      appliedFunction.current = functionKey;
      fitY.current = true;
      await refreshValues();
      if (v !== version.current) {
        return;
      }
    }
    const result = await python.call("handle_set_runs", { runs: next.map((run) => run.settings) });
    if (v !== version.current) {
      return;
    }
    setRunError(result.ok ? null : result.error.message);
    states.current = [];
    setRuns(result.ok ? next : []);
    await stepper.restart();
  }

  const functionKey = `${variables}\n${expression}`;
  const settingsKey = JSON.stringify(allSettings);
  useEffect(() => {
    if (!allValid) {
      return;
    }
    const timer = setTimeout(() => apply(functionKey, allRuns), delay.current);
    delay.current = 0;
    return () => clearTimeout(timer);
    // apply reads the latest state itself; rerun only when the settings change
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [functionKey, settingsKey]);

  // A new function, typed or from a preset (which also sets the view and start)
  function setFunction(next: string, preset?: FunctionPreset) {
    const n: Variables = detectVariables(next);
    setExpression(next);
    if (n !== variables) {
      // Runs and plot data have the old number of coordinates
      setRuns([]);
      setValues(null);
      if (!preset) {
        // Keep x, start y at 0, and show y over the same range as x
        const starts = (text: string[]) => (n === 2 ? [text[0], "0"] : [text[0]]);
        main.setStartText(starts(main.startText));
        setCompared((cs) => cs.map((c) => ({ ...c, startText: starts(c.startText) })));
        if (n === 2) {
          yRange.current = xRange.current;
        }
      }
    }
    if (preset) {
      main.setStartText(preset.start.map(String));
      setCompared((cs) => cs.map((c) => ({ ...c, startText: preset.start.map(String) })));
      xRange.current = preset.x;
      if (preset.y) {
        yRange.current = preset.y;
      }
    }
  }

  function choosePreset(preset: FunctionPreset) {
    setFunction(preset.expression, preset);
  }

  function changeExpression(text: string) {
    delay.current = TYPING_DELAY_MS;
    setFunction(text);
  }

  function startAt(editor: RunEditor, x: number, y: number) {
    editor.setStartText(variables === 1 ? [roundTo(xRange.current, x)] : [roundTo(xRange.current, x), roundTo(yRange.current, y)]);
  }

  // Somewhere in the middle 90% of the view
  function randomStart(editor: RunEditor) {
    const pick = ([a, b]: Range) => a + (0.05 + 0.9 * Math.random()) * (b - a);
    startAt(editor, pick(xRange.current), pick(yRange.current));
  }

  // The preset's start, or for a typed function the first preset's
  function defaultStart(): number[] {
    return (ALL_PRESETS.find((p) => p.expression === expression) ?? PRESETS[variables][0]).start;
  }

  function resetStart(editor: RunEditor) {
    editor.setStartText(defaultStart().map(String));
  }

  // A new comparison with default settings and the first free colour; returns its id
  function addComparison(): number {
    const used = new Set([RUN_COLORS[0].value, ...compared.map((c) => c.color)]);
    const color = RUN_COLORS.find((c) => !used.has(c.value))?.value ?? RUN_COLORS[1].value;
    const id = nextId.current++;
    setCompared([...compared, { id, color, ...newDraft(defaultStart()) }]);
    return id;
  }

  function removeComparison(id: number) {
    const rest = compared.filter((c) => c.id !== id);
    if (rest.length === 0) {
      nextId.current = 2;
    }
    setCompared(rest);
  }


  function handleViewChange(x: Range, y: Range) {
    xRange.current = x;
    yRange.current = y;
    refreshValues();
  }

  const shown = states.current.slice(0, step.iteration + 1);
  const trajectories: Trajectory[] = runs.map((run, r) => ({
    ...run,
    // A colour picked since the last rerun applies straight away
    color: compared.find((c) => c.id === run.id)?.color ?? run.color,
    path: shown.map((s) => s?.runs[r]?.point).filter((p): p is Point => !!p),
    stopped: step.state?.runs[r]?.stopped ?? null,
  }));

  return {
    pythonStatus,
    variables,
    expression,
    changeExpression,
    choosePreset,
    functionError,
    surface,
    setSurface,
    showArrows,
    setShowArrows,
    main,
    // Clicking the plot moves the main run's start
    setStartAt: (x: number, y: number) => startAt(main, x, y),
    randomStart,
    resetStart,
    runError,
    comparisons: compared.map((c, i) => ({ id: c.id, color: c.color, editor: comparisonEditors[i] })),
    canAddComparison: compared.length + 1 < MAX_RUNS,
    addComparison,
    removeComparison,
    setComparisonColor: (id: number, color: string) => setCompared((cs) => cs.map((c) => (c.id === id ? { ...c, color } : c))),
    values,
    xRange,
    yRange,
    handleViewChange,
    trajectories,
    step,
  };
}