import { useEffect, useState, useRef } from "react";
import OptimizationPlot from "./OptimizationPlot.tsx";
import Sidebar from "./Sidebar.tsx";
import ComparePanel from "./ComparePanel";
import type { BivariatePlotView, ComparisonEntry, SettingsUi } from "./types.ts";
import usePyodide from "./usePyodide.ts";

const INITIAL_SETTINGS_UI: SettingsUi = {
    mode: "Univariate",
    functionExpr: "x^2",
    algorithm: "Gradient Descent",
    x0: "0.5",
    y0: "0.5",
    learningRate: "0.1",
    momentum: "0.0",
};

const INITIAL_XLIM: [number, number] = [-1, 1];
const INITIAL_YLIM: [number, number] = [-1, 1];

const COMPARISON_COLORS = [
  "#0072B2", "#E69F00", "#009E73", "#D55E00", "#CC79A7", "#6F4EBA",
  "#A6761D", "#E7298A", "#1B9E77", "#7570B3", "#E6AB02", "#666666",
] as const;

function comparisonColor(id: number): string {
  if (id <= COMPARISON_COLORS.length) return COMPARISON_COLORS[id - 1];
  // Continue with non-repeating golden-angle hues for unusually large comparisons.
  return `hsl(${((id - 1) * 137.508) % 360}, 72%, 40%)`;
}

export default function App() {
  const [settingsUi, setSettingsUi] = useState<SettingsUi>(INITIAL_SETTINGS_UI);
  const [compareActive, setCompareActive] = useState(false);
  const [entries, setEntries] = useState<ComparisonEntry[]>([]);
  const [comparisonIteration, setComparisonIteration] = useState(0);
  const [comparisonAnimating, setComparisonAnimating] = useState(false);
  const nextId = useRef(1);
  const [isAnimating, setIsAnimating] = useState(false);
  const [showGradientField, setShowGradientField] = useState(false);
  const [bivariatePlotView, setBivariatePlotView] = useState<BivariatePlotView>("Contour heatmap");
  const xlimRef = useRef<[number, number]>(INITIAL_XLIM);
  const ylimRef = useRef<[number, number]>(INITIAL_YLIM);
  const settings = { ...settingsUi, xlim: xlimRef.current, ylim: ylimRef.current };

  function handleSettingsUiChange(newSettings: SettingsUi) {
    setIsAnimating(false);
    if (newSettings.mode !== settingsUi.mode || newSettings.functionExpr !== settingsUi.functionExpr
      || newSettings.x0 !== settingsUi.x0 || newSettings.y0 !== settingsUi.y0) {
      setComparisonAnimating(false);
      setComparisonIteration(0);
      setEntries(previous => previous.map(entry => ({ ...entry, settings: { ...entry.settings, mode: newSettings.mode, functionExpr: newSettings.functionExpr } })));
    }
    setSettingsUi(newSettings);
    api.sendInit({ ...newSettings, xlim: xlimRef.current, ylim: ylimRef.current });
  }

  function setAxisLimits(xlim: [number, number], ylim: [number, number]) {
    xlimRef.current = xlim;
    ylimRef.current = ylim;
    api.sendInit({ ...settingsUi, xlim, ylim });
  }

  function handleRandomInitialPoint() {
    const xRange = xlimRef.current[1] - xlimRef.current[0];
    const yRange = ylimRef.current[1] - ylimRef.current[0];
    const newX0 = (xlimRef.current[0] + Math.random() * xRange * 0.95).toFixed(2);
    const newY0 = (ylimRef.current[0] + Math.random() * yRange * 0.95).toFixed(2);
    const newSettings = { ...settingsUi, x0: newX0, y0: newY0 };
    handleSettingsUiChange(newSettings);
  }

  const api = usePyodide(settings);
  const { compare, comparisonBusy, isLoading } = api;
  useEffect(() => {
    if (!isLoading) compare(entries.map(entry => ({ ...entry.settings, xlim: xlimRef.current, ylim: ylimRef.current })), entries.map(entry => entry.id), comparisonIteration);
  }, [entries, comparisonIteration, compare, isLoading]);

  useEffect(() => {
    if (!comparisonAnimating || comparisonBusy || isLoading) return;
    const timeout = window.setTimeout(() => setComparisonIteration(value => value + 1), 200);
    return () => window.clearTimeout(timeout);
  }, [comparisonAnimating, comparisonBusy, isLoading, comparisonIteration]);

  const sendNextStep = api.sendNextStep;
  const currentIteration = api.currentIteration;
  const totalIterations = Math.max((api.plotData.trajectoryValues?.x.length ?? 1) - 1, 0);

  useEffect(() => {
    if (!isAnimating) return;
    const intervalId = window.setInterval(() => sendNextStep(), 200);
    return () => window.clearInterval(intervalId);
  }, [isAnimating, sendNextStep]);

  //if (api.isLoading) {
    //return <LoadingScreen message="Loading" />;
  //}

  return (
		<div className="grid grid-cols-[1fr_12fr_1fr] lg:h-dvh bg-stone-50">
			<div/>
			<div className="flex flex-col gap-6 p-6 min-h-0">
				<header className="grid grid-cols-1 items-center gap-3 lg:grid-cols-[2fr_1fr] lg:gap-12">
					<div className="title">Optimization Trajectory Visualizer</div>
					{/* <div className="desc"></div> */}
					<div
						id="runtime-status"
						className={`status translate-y-4 w-fit justify-self-end ${api.isLoading ? "loading" : "ready"}`}
						role="status"
						aria-live="polite"
					>
						<span></span>
						{api.isLoading ? "Loading Python runtime…" : "Python runtime ready"}
					</div>
				</header>

				<div className="grid grid-cols-1 lg:grid-cols-[2fr_1fr] gap-12 flex-1 min-h-0 overflow-auto">
					<div className="flex h-full min-h-0 min-w-0 flex-col gap-2">
            {settingsUi.mode === "Bivariate" && (
              <div className="flex w-full border-b border-gray-200" role="tablist" aria-label="Bivariate plot view">
                {(["Contour heatmap", "Function surface"] as const).map((view) => (
                  <button
                    key={view}
                    type="button"
                    role="tab"
                    aria-selected={bivariatePlotView === view}
                    onClick={() => setBivariatePlotView(view)}
                    className={`w-1/2 cursor-pointer rounded px-5 py-1 ${bivariatePlotView === view
                      ? "border-b-2 border-orange-400 text-orange-400"
                      : "text-gray-950 hover:bg-gray-200"}`}
                  >
                    {view}
                  </button>
                ))}
              </div>
            )}
            <OptimizationPlot 
              data={compareActive ? { functionValues: api.plotData.functionValues } : api.plotData}
              comparisons={compareActive ? entries.map(entry => ({ ...entry, values: api.comparison.find(item => item.id === entry.id)?.values })) : undefined}
              isLoading={api.isLoading}
              xlim={xlimRef.current}
              ylim={ylimRef.current}
              setAxisLimits={setAxisLimits}
              currentIteration={currentIteration}
              showGradientField={showGradientField}
              bivariatePlotView={bivariatePlotView}
            />
					</div>
					<Sidebar
            onTabChange={(tab) => { setCompareActive(tab === "Compare"); setIsAnimating(false); setComparisonAnimating(false); }}
            comparison={<ComparePanel settings={settingsUi} entries={entries} trajectories={api.comparison}
              iteration={comparisonIteration} busy={comparisonBusy} ready={!api.isLoading} animating={comparisonAnimating}
              onAdd={(configured) => {
                const id = nextId.current++;
                setComparisonIteration(0);
                setEntries(previous => [...previous, {
                  id,
                  settings: configured,
                  color: comparisonColor(id),
                  visible: true,
                }]);
              }}
              onRemove={(id) => { setComparisonAnimating(false); setEntries(previous => previous.filter(entry => entry.id !== id)); }}
              onToggleVisibility={(id) => setEntries(previous => previous.map(entry => (
                entry.id === id ? { ...entry, visible: !entry.visible } : entry
              )))}
              onAction={(action) => {
                if (action === "Animate") setComparisonAnimating(true);
                else if (action === "Pause") setComparisonAnimating(false);
                else if (action === "Reset") setComparisonIteration(0);
                else setComparisonIteration(value => Math.max(0, value + (action === "Next" ? 1 : -1)));
              }}
            />}
						settings={settingsUi} 
						setSettings={handleSettingsUiChange} 
						onRandomInitialPoint={handleRandomInitialPoint}
						trajectoryValues={api.plotData.trajectoryValues}

						onReset={() => {
							setIsAnimating(false);
							api.sendReset();
						}}
						onNextStep={api.sendNextStep}
						onPrevStep={api.sendPrevStep}
						onAnimate={() => setIsAnimating(true)}
						onPause={() => setIsAnimating(false)}
						isAnimating={isAnimating}
						currentIteration={currentIteration}
						totalIterations={totalIterations}
						showGradientField={showGradientField}
						onShowGradientFieldChange={setShowGradientField}
					/>
				</div>
    </div>
		<div/>
	</div>
  );
}