import { useRef } from "react";
import Plot from "react-plotly.js";
import type { PlotData } from "./types";

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

export default function OptimizationPlot({ data, 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 plotRef = useRef<any>(null);
  const layoutRef = useRef<any>({
    dragmode: 'pan',
    showlegend: false,
    xaxis: {
      title: { text: 'x' },
      range: xlim,
    },
    yaxis: {
      title: { text: 'y' },
      range: ylim,
    }
  })

  if (z.length === 0) {
    return (
      <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}
        className="w-full h-full"
        config={{ 
          responsive: true,
          displayModeBar: false,
          scrollZoom: true,
        }}
      />
    );
  } else {
    return (
      <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,
            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 h-full"
        config={{ 
          responsive: true,
          displayModeBar: false,
          scrollZoom: true,
        }}
      />
    ); 
  }
}