import { useMemo, useRef } from "react";
import Plot from "react-plotly.js";
import type { BivariatePlotView, ComparisonEntry, TrajectoryValues, PlotData } from "./types";
import { Card } from "@elvis/ui";

interface OptimizationPlotProps {
  comparisons?: (ComparisonEntry & { values?: TrajectoryValues })[];
  data: PlotData;
  isLoading: boolean;
  xlim: [number, number];
  ylim: [number, number];
  setAxisLimits: (xlim: [number, number], ylim: [number, number]) => void;
  currentIteration: number;
  showGradientField: boolean;
  bivariatePlotView: BivariatePlotView;
}

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

  const isSurface = bivariatePlotView === "Function surface" && Boolean(data.functionValues?.z);
  const comparisonTraces = (comparisons ?? []).flatMap((entry, index) => entry.visible ? [
    { x: entry.values?.x ?? [], y: entry.values?.y ?? [], z: isSurface ? entry.values?.z ?? [] : undefined,
      type: isSurface ? "scatter3d" as const : "scatter" as const, mode: "lines+markers" as const,
      name: `${index + 1}. ${entry.settings.algorithm}`, legendgroup: String(entry.id), showlegend: true,
      line: { color: entry.color, width: isSurface ? 6 : 2 },
      marker: { color: entry.color, size: isSurface ? 2.5 : 6 } },
    { x: entry.values?.x.slice(-1) ?? [], y: entry.values?.y.slice(-1) ?? [], z: isSurface ? entry.values?.z?.slice(-1) ?? [] : undefined,
      type: isSurface ? "scatter3d" as const : "scatter" as const, mode: "markers" as const,
      name: `${index + 1}. ${entry.settings.algorithm}`, legendgroup: String(entry.id), showlegend: false,
      marker: { color: entry.color, size: isSurface ? 4 : 12, line: { color: "white", width: isSurface ? 1 : 2 } } }
  ] : []);
  // 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 gradientX: number[][] = data.functionValues?.gradientX ?? [];
  let gradientY: number[][] = data.functionValues?.gradientY ?? [];

	console.log('gradientX', gradientX)
	console.log('gradientY', gradientY)


  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 = useMemo<[number, number] | null>(() => {
    if (z.length === 0) return null;

    let zMin = Infinity;
    let zMax = -Infinity;
    for (const row of z) {
      for (const value of row) {
        if (!Number.isFinite(value)) continue;
        zMin = Math.min(zMin, value);
        zMax = Math.max(zMax, value);
      }
    }

    if (!Number.isFinite(zMin) || !Number.isFinite(zMax)) return null;
    const padding = zMax === zMin ? Math.max(Math.abs(zMin) * 0.1, 1) : (zMax - zMin) * 0.1;
    return [zMin - padding, zMax + padding];
  }, [z]);

  const gradientField = useMemo(() => {
    const emptyField = {
      shaftX: [] as (number | null)[],
      shaftY: [] as (number | null)[],
      headX: [] as number[],
      headY: [] as number[],
      headAngles: [] as number[],
      hoverData: [] as number[][],
    };
    if (
      !showGradientField
      || gradientX.length !== y.length
      || gradientY.length !== y.length
      || x.length === 0
      || y.length === 0
    ) return emptyField;

    const rowStep = 5 //Math.max(1, Math.ceil(y.length / 10));
    const columnStep = 5 //Math.max(1, Math.ceil(x.length / 10));
    const shaftX: (number | null)[] = [];
    const shaftY: (number | null)[] = [];
    const headX: number[] = [];
    const headY: number[] = [];
    const headAngles: number[] = [];
    const hoverData: number[][] = [];

		const gradientLengths: number[] = [];

		for (let row = 0; row < gradientX.length; row++) {
			for (let column = 0; column < gradientX[row].length; column++) {
				const gx = gradientX[row][column];
				const gy = gradientY[row][column];

				const length = Math.hypot(gx, gy);

				if (Number.isFinite(length)) {
					gradientLengths.push(length);
				}
			}
		}

		const averageGradientLength = gradientLengths.length > 0 ? 
			gradientLengths.reduce((sum, length) => sum + length, 0) / gradientLengths.length
				: 0;
				
		console.log('average_grad_length', averageGradientLength);
		console.log('cell diagonal', Math.hypot(x[columnStep]-x[0], y[rowStep]-y[0]));
		const grad_scale = 1.8*averageGradientLength / Math.hypot(x[columnStep]-x[0], y[rowStep]-y[0]);
		console.log('grad_scale', grad_scale);

    for (let row = 0; row < y.length; row += rowStep) {
      if (!gradientX[row] || !gradientY[row]) continue;
      for (let column = 0; column < x.length; column += columnStep) {
        const gx = gradientX[row][column];
        const gy = gradientY[row][column];
        const magnitude = Math.hypot(gx, gy);
        if (!Number.isFinite(magnitude) || magnitude === 0) continue;

        const vx = gx / grad_scale;
        const vy = gy / grad_scale;
        const arrowX = x[column] + vx;
        const arrowY = y[row] + vy;
        shaftX.push(x[column], arrowX, null);
        shaftY.push(y[row], arrowY, null);
        headX.push(arrowX);
        headY.push(arrowY);
        // Plotly marker angles rotate clockwise from the upward direction.
        headAngles.push(Math.atan2(vx, vy) * 180 / Math.PI);
        hoverData.push([x[column], y[row], gx, gy]);
      }
    }

    return { shaftX, shaftY, headX, headY, headAngles, hoverData };
  }, [showGradientField, x, y, gradientX, gradientY]);

  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;

  function zoomSurfaceAxes(event: React.WheelEvent) {
    event.preventDefault();
    const factor = event.deltaY > 0 ? 1.2 : 1 / 1.2;
    const resizeRange = (range: [number, number]): [number, number] => {
      const center = (range[0] + range[1]) / 2;
      const halfWidth = (range[1] - range[0]) * factor / 2;
      return [center - halfWidth, center + halfWidth];
    };
    setAxisLimits(resizeRange(xlim), resizeRange(ylim));
  }

  if (isLoading && hasNoData) {
    return (
      <Card className="min-h-[320px] flex flex-1 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] flex-1">
        <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",
              showlegend: false,
            },
            {
              x: trajX,
              y: trajY,
              type: 'scatter',
              mode: 'lines+markers',
              line: { color: '#d97871', width: 2 },
              marker: { color: '#d97871', size: 10 },
              hoverinfo: "skip",
              showlegend: false,
            },
            {
              x: trajX[currentIteration] !== undefined ? [trajX[currentIteration]] : [],
              y: trajY[currentIteration] !== undefined ? [trajY[currentIteration]] : [],
              type: 'scatter',
              mode: 'markers',
              marker: { color: 'red', size: 12 },
              hoverinfo: "skip",
              showlegend: false,
            }
            ,...comparisonTraces
          ]}
          layout={{ ...layoutRef.current, showlegend: comparisons !== undefined, legend: { x: 0.02, y: 0.98, xanchor: "left", yanchor: "top", bgcolor: "rgba(255,255,255,0.75)" } }}
          style={{ width: '100%', height: '100%' }}
          config={{ 
            responsive: true,
            displayModeBar: true,
            scrollZoom: true,
          }}
        />
      </Card>
    );
  } else {
    if (isSurface) {
      return (
        <div className="min-h-[420px] flex-1" onWheel={zoomSurfaceAxes}>
        <Card className="h-full">
          <Plot
            ref={plotRef}
            data={[
              {
                x, y, z,
                type: 'surface',
                colorscale: 'Viridis',
                opacity: 0.82,
                showscale: true,
                hovertemplate: 'x: %{x:.5g}<br>y: %{y:.5g}<br>z: %{z:.5g}<extra></extra>',
              },
              {
                x: trajX, y: trajY, z: trajZ,
                type: 'scatter3d', mode: 'lines+markers',
                line: { color: '#d97871', width: 6 },
                marker: { color: '#d97871', size: 2.5 },
                hoverinfo: 'skip', showlegend: false,
              },
              {
                x: trajX[currentIteration] !== undefined ? [trajX[currentIteration]] : [],
                y: trajY[currentIteration] !== undefined ? [trajY[currentIteration]] : [],
                z: trajZ[currentIteration] !== undefined ? [trajZ[currentIteration]] : [],
                type: 'scatter3d', mode: 'markers',
                marker: { color: 'red', size: 7 },
                hoverinfo: 'skip', showlegend: false,
              },
              ...comparisonTraces,
            ]}
            layout={{
              // Keep the user's camera angle while animation updates the trajectory data.
              uirevision: 'bivariate-surface',
              dragmode: 'orbit',
              showlegend: comparisons !== undefined,
              legend: { x: 0.02, y: 0.98, xanchor: 'left', yanchor: 'top', bgcolor: 'rgba(255,255,255,0.75)' },
              scene: {
                aspectmode: 'cube',
                xaxis: { title: { text: 'x' }, range: xlim, autorange: false },
                yaxis: { title: { text: 'y' }, range: ylim, autorange: false },
                zaxis: { title: { text: 'z' }, range: colorScaleRange ?? undefined, autorange: false },
              },
              margin: { t: 30, r: 20, b: 20, l: 20 },
            }}
            style={{ width: '100%', height: '100%' }}
            config={{ responsive: true, displayModeBar: true, scrollZoom: false }}
          />
        </Card>
        </div>
      );
    }

    return (
      <Card className="min-h-[420px] flex-1">
        <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",
              showlegend: false,
              contours: {
                coloring: "heatmap",
                showlines: false,
              }
            },
            {
              x: gradientField.shaftX,
              y: gradientField.shaftY,
              type: 'scatter',
              mode: 'lines',
              line: { color: 'rgba(255, 255, 255, 0.8)', width: 1 },
              hoverinfo: 'skip',
              showlegend: false,
              cliponaxis: true,
            },
            {
              x: gradientField.headX,
              y: gradientField.headY,
              customdata: gradientField.hoverData,
              type: 'scatter',
              mode: 'markers',
              marker: {
                symbol: 'triangle-up',
                size: 5,
                color: 'rgba(255, 255, 255, 0.8)',
                angle: gradientField.headAngles,
                angleref: 'up',
              } as any,
              hovertemplate:
                'x: %{customdata[0]:.5g}<br>'
                + 'y: %{customdata[1]:.5g}<br>'
                + 'gx: %{customdata[2]:.5g}<br>'
                + 'gy: %{customdata[3]:.5g}<extra></extra>',
              showlegend: false,
              cliponaxis: true,
            },
            {
              x: trajX,
              y: trajY,
              customdata: trajZ,
              type: 'scatter',
              mode: 'lines+markers',
              line: { color: '#d97871', width: 2 },
              marker: { color: '#d97871', size: 10 },
              hovertemplate: 'x: %{x:.5g}<br>y: %{y:.5g}<br>f(x, y): %{customdata:.5g}<extra></extra>',
              showlegend: false,
            },
            {
              x: trajX[currentIteration] !== undefined ? [trajX[currentIteration]] : [],
              y: trajY[currentIteration] !== undefined ? [trajY[currentIteration]] : [],
              customdata: trajZ[currentIteration] !== undefined ? [trajZ[currentIteration]] : [],
              type: 'scatter',
              mode: 'markers',
              marker: { color: 'red', size: 12 },
              hovertemplate: 'x: %{x:.5g}<br>y: %{y:.5g}<br>f(x, y): %{customdata:.5g}<extra></extra>',
              showlegend: false,
            }
            ,...comparisonTraces
          ]}
          layout={{ ...layoutRef.current, annotations: [], showlegend: comparisons !== undefined, legend: { x: 0.02, y: 0.98, xanchor: "left", yanchor: "top", bgcolor: "rgba(255,255,255,0.75)" } }}
          style={{ width: '100%', height: '100%' }}
          config={{ 
            responsive: true,
            displayModeBar: true,
            scrollZoom: true,
          }}
        />
      </Card>
    ); 
  }
}