import { useState } from "react";
import { ChevronDown, X } from "lucide-react";
import {
  AppLayout,
  Button,
  ColorPicker,
  Divider,
  Dropdown,
  IconButton,
  InputField,
  InputSection,
  ResultSection,
  SegmentedControl,
  StepControls,
  type ResultItem,
} from "@elvis/ui";

import { algorithmById } from "./algorithms";
import { RUN_COLORS } from "./colors";
import { ALL_PRESETS, CUSTOM, PRESETS } from "./functions";
import HistoryPlot from "./HistoryPlot";
import OptimisationPlot from "./OptimisationPlot";
import OptimiserControls from "./OptimiserControls";
import Usage from "./Usage";
import useOptimisation, { type RunEditor, type Trajectory } from "./useOptimisation";
import type { Point } from "./types";

const PRESET_GROUPS = [
  { label: "Univariate: f(x)", options: PRESETS[1].map((p) => p.name) },
  { label: "Bivariate: f(x, y)", options: PRESETS[2].map((p) => p.name) },
];

const PLOT_OPTIONS = [
  { value: "contour", label: "Contour" },
  { value: "surface", label: "Surface" },
] as const;

const helpClass = "text-sm text-ink-muted";

export default function App() {
  const app = useOptimisation();
  const preset = ALL_PRESETS.find((p) => p.expression === app.expression);
  const twoVariables = app.variables === 2;
  const main = app.trajectories[0];
  const point = main?.path[main.path.length - 1];
  // The comparison card being edited; the others are collapsed
  const [openCard, setOpenCard] = useState<number | null>(null);

  return (
    <AppLayout
      title="Optimisation"
      usage={<Usage />}
      examples={<p>Examples are coming soon.</p>}
      tabs={[{ label: "Configure", content: <>
          <InputSection label="Function">
            <Dropdown
              label="Function"
              hideLabel
              // "Custom" is listed only while the box holds a function of the student's own
              options={preset ? [] : [CUSTOM]}
              groupedOptions={PRESET_GROUPS}
              activeOption={preset?.name ?? CUSTOM}
              onChange={(name) => {
                const chosen = ALL_PRESETS.find((p) => p.name === name);
                if (chosen) {
                  app.choosePreset(chosen);
                }
              }}
            />
            <InputField
              label={twoVariables ? "f(x, y) =" : "f(x) ="}
              hideLabel
              prefix={twoVariables ? "f(x, y) =" : "f(x) ="}
              value={app.expression}
              onChange={app.changeExpression}
              error={app.functionError}
              mono
              rows={2}
            />
            {twoVariables && (
              // The plot view in two thirds, the arrows toggle (contour only) in the last
              <div className="grid grid-cols-3 items-center">
                <div className="col-span-2">
                  <SegmentedControl
                    label="Plot"
                    options={PLOT_OPTIONS}
                    value={app.surface ? "surface" : "contour"}
                    onChange={(v) => app.setSurface(v === "surface")}
                  />
                </div>
                <label
                  title={app.surface ? "Only on the contour plot" : undefined}
                  className={`flex items-center justify-end gap-1.5 whitespace-nowrap text-sm text-gray-700 ${app.surface ? "opacity-50" : "cursor-pointer"}`}
                >
                  <input
                    type="checkbox"
                    checked={app.showArrows}
                    disabled={app.surface}
                    onChange={(e) => app.setShowArrows(e.target.checked)}
                    className="size-4 accent-ink"
                  />
                  Gradient field
                </label>
              </div>
            )}
          </InputSection>
          <Divider />
          <InputSection label="Optimiser">
            <OptimiserControls
              editor={app.main}
              twoVariables={twoVariables}
              onResetStart={() => app.resetStart(app.main)}
              onRandomStart={() => app.randomStart(app.main)}
              hint={twoVariables && app.surface ? undefined : "Or click the plot to move the initial point."}
            />
          </InputSection>
        </> },
        { label: "Compare", content: <>
          <p className={helpClass}>Add optimisers to compare trajectories</p>
          <ul className="flex flex-col gap-2">
            {/* The main run: the same settings as Configure, with a fixed colour */}
            <RunCard
              id={1}
              color={RUN_COLORS[0].value}
              editor={app.main}
              trajectory={main}
              note="Also set in Configure"
              open={openCard === 1}
              onToggle={() => setOpenCard(openCard === 1 ? null : 1)}
              twoVariables={twoVariables}
              onResetStart={() => app.resetStart(app.main)}
              onRandomStart={() => app.randomStart(app.main)}
            />
            {app.comparisons.map((c) => (
              <RunCard
                key={c.id}
                id={c.id}
                color={c.color}
                editor={c.editor}
                trajectory={app.trajectories.find((t) => t.id === c.id)}
                taken={[RUN_COLORS[0].value, ...app.comparisons.map((o) => o.color)]}
                open={openCard === c.id}
                onToggle={() => setOpenCard(openCard === c.id ? null : c.id)}
                onColorChange={(color) => app.setComparisonColor(c.id, color)}
                onRemove={() => app.removeComparison(c.id)}
                twoVariables={twoVariables}
                onResetStart={() => app.resetStart(c.editor)}
                onRandomStart={() => app.randomStart(c.editor)}
              />
            ))}
          </ul>
          <Button label="+ Add optimiser" onClick={() => setOpenCard(app.addComparison())} disabled={!app.canAddComparison} />
        </> },
        { label: "Metrics", content: point ? <>
          <ResultSection title="Current point" items={pointItems(point)} />
          <Divider />
          <InputSection label="History">
            <HistoryPlot trajectories={app.trajectories} />
          </InputSection>
        </> : <p className={helpClass}>Step the optimiser to see the current point and its history.</p> },
      ]}
      footer={<>
        {app.pythonStatus === "loading" && (
          <p className={helpClass}>Loading Python… actions will run once it is ready.</p>
        )}
        {app.pythonStatus === "failed" && (
          <p className="text-red-600 font-semibold">Python failed to load. Try reloading the page.</p>
        )}
        {app.runError && <p className="text-red-600 font-semibold">{app.runError}</p>}
        <StepControls stepper={app.step} />
        {app.trajectories.filter((t) => t.stopped).map((t) => (
          <p key={t.label} className={helpClass}>
            {t.label} stopped at iteration {t.path.length - 1}. {t.stopped}
          </p>
        ))}
      </>}
    >
      <OptimisationPlot
        values={app.values}
        trajectories={app.trajectories}
        xRange={app.xRange}
        yRange={app.yRange}
        onViewChange={app.handleViewChange}
        onClick={app.setStartAt}
        surface={app.surface}
        arrows={app.showArrows}
      />
    </AppLayout>
  );
}

// f and its derivatives at the point, as the panel lists them
function pointItems(p: Point): ResultItem[] {
  if (p.position.length === 1) {
    return [
      { name: "x", kind: "scalar", value: p.position[0] },
      { name: "f(x)", kind: "scalar", value: p.value },
      { name: "f′(x)", kind: "scalar", value: p.gradient[0] },
      { name: "f″(x)", kind: "scalar", value: p.hessian[0][0] },
    ];
  }
  return [
    { name: "Position", kind: "vector", value: { labels: ["x", "y"], values: p.position } },
    { name: "f(x, y)", kind: "scalar", value: p.value },
    { name: "Gradient ∇f", kind: "vector", value: { labels: ["∂f/∂x", "∂f/∂y"], values: p.gradient } },
    { name: "Hessian", kind: "matrix", value: { rows: ["x", "y"], columns: ["x", "y"], values: p.hessian } },
  ];
}

// f at a run's latest point, if it has one
function CurrentValue({ trajectory }: { trajectory?: Trajectory }) {
  const last = trajectory?.path[trajectory.path.length - 1];
  return last ? <span className="shrink-0 font-mono text-ink">f = {Number(last.value.toPrecision(3))}</span> : null;
}

interface RunCardProps {
  id: number;
  color: string;
  editor: RunEditor;
  trajectory?: Trajectory;
  // A line under the name, e.g. that the main run is also set in Configure
  note?: string;
  open: boolean;
  onToggle: () => void;
  // Without these (the main run) the colour is fixed and the run can't be removed
  onColorChange?: (color: string) => void;
  onRemove?: () => void;
  // Colours used by the other runs
  taken?: string[];
  twoVariables: boolean;
  onResetStart: () => void;
  onRandomStart: () => void;
}

// A run in the Compare tab: one line with its colour, optimiser and f, which
// opens to the same optimiser controls as Configure
function RunCard({ id, color, editor, trajectory, note, open, onToggle, onColorChange, onRemove, taken = [], twoVariables, onResetStart, onRandomStart }: RunCardProps) {
  const label = `${id} · ${algorithmById(editor.algorithm).label}`;
  return (
    <li className="rounded-md border-[1.5px] border-grid">
      <div className="flex items-center gap-2 py-1 pl-1 pr-1">
        {onColorChange ? (
          <ColorPicker label={`Colour of run ${id}`} colors={RUN_COLORS} value={color} onChange={onColorChange} taken={taken} />
        ) : (
          <span className="flex size-9 shrink-0 items-center justify-center" title={`Run ${id}'s colour is fixed`}>
            <span aria-hidden className="block size-6 rounded-full border-2 border-white ring-[1.5px] ring-grid" style={{ backgroundColor: color }} />
          </span>
        )}
        <button
          type="button"
          onClick={onToggle}
          aria-expanded={open}
          className="flex min-w-0 flex-1 cursor-pointer items-center gap-2 self-stretch rounded px-1 text-left text-sm text-gray-700 hover:bg-paper"
        >
          <span className="min-w-0 flex-1">
            <span className="block truncate">{label}</span>
            {note && <span className="block text-xs text-ink-muted">{note}</span>}
            {!editor.startValid && <span className="block text-xs font-semibold text-red-600">Check the initial point</span>}
          </span>
          <CurrentValue trajectory={trajectory} />
          <ChevronDown size={16} aria-hidden className={`shrink-0 text-ink-muted transition-transform ${open ? "rotate-180" : ""}`} />
        </button>
        {onRemove && (
          <IconButton label={`Remove run ${id}`} onClick={onRemove}>
            <X size={16} />
          </IconButton>
        )}
      </div>
      {open && (
        <div className="flex flex-col gap-4 border-t border-grid p-3">
          <OptimiserControls editor={editor} twoVariables={twoVariables} onResetStart={onResetStart} onRandomStart={onRandomStart} />
        </div>
      )}
    </li>
  );
}