frontends/react/src/OptimizationPlot.tsx

import { useState, useRef, useEffect } from "react";
import Plot from "react-plotly.js";
import type { PlotData } from "./types";
import { Button, Card } from "@elvis/ui";

interface OptimizationPlotProps {
  data: PlotData;
  isLoading: boolean;
  xlim: [number, number];
  ylim: [number, number];
  setAxisLimits: (xlim: [number, number], ylim: [number, number]) => void;
}

export default function OptimizationPlot({ data, isLoading, xlim, ylim, setAxisLimits }: OptimizationPlotProps) {

  // data
  let x: number[] = data.functionValues ? data.functionValues.x : [];
  let y: number[] = data.functionValues ? data.functionValues.y : [];
  let z: number[][] = data.functionValues && data.functionValues.z ? data.functionValues.z : [];

  let trajX: number[] = data.trajectoryValues ? data.trajectoryValues.x : [];
  let trajY: number[] = data.trajectoryValues ? data.trajectoryValues.y : [];
  let trajZ: number[] = data.trajectoryValues && data.trajectoryValues.z ? data.trajectoryValues.z : [];
  
  const [colorScaleRange, setColorScaleRange] = useState<[number, number] | null>(null);
  const nextColorScaleRangeRef = useRef<[number, number] | null>(null);

  useEffect(() => {
    if (z) {
      if (!z || z.length === 0) {
        nextColorScaleRangeRef.current = null;
        return;
      }

      let zMin = Infinity;
      let zMax = -Infinity;

      for (let i = 0; i < z.length; i++) {
        for (let j = 0; j < z[i].length; j++) {
          const v = z[i][j];
          if (!Number.isFinite(v)) {
            continue;
          }
          if (v < zMin) {
            zMin = v;
          }
          if (v > zMax) {
            zMax = v;
          }
        }
      }

      const padding = (zMax - zMin) * 0.1;
      nextColorScaleRangeRef.current = [zMin - padding, zMax + padding];

      if (colorScaleRange === null) {
        setColorScaleRange(nextColorScaleRangeRef.current);
      }
    }
  }, [z]);

  function updateColorScaleRange() {
    if (nextColorScaleRangeRef.current) {
      setColorScaleRange(nextColorScaleRangeRef.current);
      nextColorScaleRangeRef.current = null;
    }
  }

  const plotRef = useRef<any>(null);
  const layoutRef = useRef<any>({
    dragmode: 'pan',
    showlegend: false,
    xaxis: {
      title: { text: 'x' },
      range: xlim,
    },
    yaxis: {
      title: { text: 'y' },
      range: ylim,
    },
    margin: { t: 40, r: 40, b: 40, l: 40 }
  })

  const hasNoData = !data.functionValues && !data.trajectoryValues;

  if (isLoading && hasNoData) {
    return (
      <Card className="min-h-[320px] flex items-center justify-center p-6">
        <div className="text-center">
          <div className="text-lg font-medium">Loading...</div>
        </div>
      </Card>
    );
  }

  if (z.length === 0) {
    return (
      <Card className="min-h-[320px]">
        <Plot
          ref={plotRef}
          onRelayout={(event) => {
            const x0 = event['xaxis.range[0]'];
            const x1 = event['xaxis.range[1]'];
            const y0 = event['yaxis.range[0]'];
            const y1 = event['yaxis.range[1]'];

            if (
              typeof x0 === "number"
              && typeof x1 === "number"
              && typeof y0 === "number"
              && typeof y1 === "number"
            ) {
              setAxisLimits([x0, x1], [y0, y1]);
            }
          }}
          data={[
            {
              x: x,
              y: y,
              type: 'scatter',
              mode: 'lines',
              line: { color: '#1f77b4', width: 2 },
              hoverinfo: "skip",
            },
            {
              x: trajX,
              y: trajY,
              type: 'scatter',
              mode: 'lines+markers',
              line: { color: '#d97871', width: 2 },
              marker: { color: '#d97871', size: 10 },
              hoverinfo: "skip",
            },
            {
              x: trajX.length > 0 ? [trajX.at(-1)!] : [],
              y: trajY.length > 0 ? [trajY.at(-1)!] : [],
              type: 'scatter',
              mode: 'markers',
              marker: { color: 'red', size: 12 },
              hoverinfo: "skip",
            }
          ]}
          layout={layoutRef.current}
          style={{ width: '100%', height: '100%' }}
          config={{ 
            responsive: true,
            displayModeBar: true,
            scrollZoom: true,
          }}
        />
      </Card>
    );
  } else {
    return (
      <Card className="flex flex-col min-h-[420px]">
        <Plot
          ref={plotRef}
          onRelayout={(event) => {
            const x0 = event['xaxis.range[0]'];
            const x1 = event['xaxis.range[1]'];
            const y0 = event['yaxis.range[0]'];
            const y1 = event['yaxis.range[1]'];

            if (
              typeof x0 === "number"
              && typeof x1 === "number"
              && typeof y0 === "number"
              && typeof y1 === "number"
            ) {
              setAxisLimits([x0, x1], [y0, y1]);
            }
          }}
          data={[
            {
              x: x,
              y: y,
              z: z,
              zmin: colorScaleRange?.[0],
              zmax: colorScaleRange?.[1],
              type: 'contour',
              colorscale: 'Viridis',
              hoverinfo: "skip",
              contours: {
                coloring: "heatmap",
                showlines: false,
              }
            },
            {
              x: trajX,
              y: trajY,
              z: trajZ,
              type: 'scatter',
              mode: 'lines+markers',
              line: { color: '#d97871', width: 2 },
              marker: { color: '#d97871', size: 10 },
              hoverinfo: "skip",
            },
            {
              x: trajX.length > 0 ? [trajX.at(-1)!] : [],
              y: trajY.length > 0 ? [trajY.at(-1)!] : [],
              z: trajZ.length > 0 ? [trajZ.at(-1)!] : [],
              type: 'scatter',
              mode: 'markers',
              marker: { color: 'red', size: 12 },
              hoverinfo: "skip",
            }
          ]}
          layout={layoutRef.current}
          className="w-full flex-1"
          config={{ 
            responsive: true,
            displayModeBar: false,
            scrollZoom: true,
          }}
        />
        <div className="mt-2 flex justify-end">
          <Button label="Update Color Scale" onClick={updateColorScaleRange} />
        </div>
      </Card>
    ); 
  }
}