import type React from "react";
import { useEffect, useRef, useState } from "react";
import * as papa from "papaparse";
import type { Route } from "./+types/home";

export function meta({ }: Route.MetaArgs) {
  return [
    { title: "3dplot-gen" },
    { name: "description", content: "CSV → 3D Plot generator" },
  ];
}

type Row = Record<string, number>;

type PlotType = "scatter" | "surface";

type ThemeScheme = "light" | "dark" | "transparent";

const THEME_SCHEMES: Record<
  ThemeScheme,
  {
    paperBg: string;
    plotBg: string;
    fontColor: string;
    gridColor: string;
    mainClass: string;
  }
> = {
  dark: {
    paperBg: "#0b0d12",
    plotBg: "#0b0d12",
    fontColor: "#cfd3ff",
    gridColor: "rgba(148, 163, 184, 0.25)",
    mainClass: "bg-[#0b0d12]",
  },
  light: {
    paperBg: "#f8fafc",
    plotBg: "#f8fafc",
    fontColor: "#1e293b",
    gridColor: "rgba(148, 163, 184, 0.5)",
    mainClass: "bg-slate-100",
  },
  transparent: {
    paperBg: "rgba(0,0,0,0)",
    plotBg: "rgba(0,0,0,0)",
    fontColor: "#475569",
    gridColor: "rgba(100, 116, 139, 0.4)",
    mainClass: "bg-transparent",
  },
};

const COLORSCALES = [
  "Viridis",
  "Plasma",
  "Inferno",
  "Magma",
  "Cividis",
  "Hot",
  "Electric",
  "Bluered",
  "Portland",
  "Rainbow",
  "RdBu",
  "Earth",
  "Teal",
  "Turbo",
] as const;

export default function Home() {
  const plotRef = useRef<HTMLDivElement>(null);

  const [rows, setRows] = useState<Row[]>([]);
  const [columns, setColumns] = useState<string[]>([]);
  const [x, setX] = useState("");
  const [y, setY] = useState("");
  const [z, setZ] = useState("");
  const [plotType, setPlotType] = useState<PlotType>("scatter");
  const [colorscale, setColorscale] = useState<string>(COLORSCALES[0]);
  const [themeScheme, setThemeScheme] = useState<ThemeScheme>("dark");

  // First row after header = axis sizes (X max, Y max, Z max); rest = plot data.
  const plotData = rows.length > 1 ? rows.slice(1) : rows;
  const axisSizes =
    rows.length > 1 && x && y && z
      ? (() => {
        const sx = Number(rows[0][x]);
        const sy = Number(rows[0][y]);
        const sz = Number(rows[0][z]);
        return [sx, sy, sz].every((n) => Number.isFinite(n))
          ? { x: sx, y: sy, z: sz }
          : null;
      })()
      : null;

  // Build surface grid from scattered (x,y,z) points. Returns null if data isn't a grid.
  function buildSurfaceGrid(): {
    xUnique: number[];
    yUnique: number[];
    zMatrix: number[][];
  } | null {
    const xVals = plotData.map((r) => Number(r[x]));
    const yVals = plotData.map((r) => Number(r[y]));
    const zVals = plotData.map((r) => Number(r[z]));
    const xUnique = [...new Set(xVals)].sort((a, b) => a - b);
    const yUnique = [...new Set(yVals)].sort((a, b) => a - b);
    const keyToZ = new Map<string, number>();
    for (let i = 0; i < plotData.length; i++) {
      const k = `${xVals[i]},${yVals[i]}`;
      if (!keyToZ.has(k)) keyToZ.set(k, zVals[i]);
    }
    if (xUnique.length * yUnique.length !== keyToZ.size) return null;
    const zMatrix: number[][] = yUnique.map((yi) =>
      xUnique.map((xj) => keyToZ.get(`${xj},${yi}`) ?? NaN)
    );
    return { xUnique, yUnique, zMatrix };
  }

  // Render plot when axes or options change
  useEffect(() => {
    if (!plotRef.current || !x || !y || !z) return;

    let cancelled = false;

    (async () => {
      const Plotly = await import("plotly.js-dist-min");
      if (cancelled) return;

      const xVals = plotData.map((r) => Number(r[x]));
      const yVals = plotData.map((r) => Number(r[y]));
      const zVals = plotData.map((r) => Number(r[z]));

      const grid = buildSurfaceGrid();
      const useSurface = plotType === "surface" && grid !== null;

      const trace = useSurface && grid
        ? {
          type: "surface" as const,
          x: grid.xUnique,
          y: grid.yUnique,
          z: grid.zMatrix,
          colorscale: colorscale as "Viridis",
          colorbar: { title: { text: z } },
        }
        : {
          type: "scatter3d" as const,
          mode: "markers" as const,
          x: xVals,
          y: yVals,
          z: zVals,
          marker: {
            size: 4,
            color: zVals,
            colorscale: colorscale as "Viridis",
          },
        };

      const theme = THEME_SCHEMES[themeScheme];
      const sceneLayout: {
        bgcolor?: string;
        xaxis: {
          title: { text: string };
          range?: [number, number];
          gridcolor?: string;
          tickfont?: { color: string };
        };
        yaxis: {
          title: { text: string };
          range?: [number, number];
          gridcolor?: string;
          tickfont?: { color: string };
        };
        zaxis: {
          title: { text: string };
          range?: [number, number];
          gridcolor?: string;
          tickfont?: { color: string };
        };
      } = {
        bgcolor: theme.plotBg,
        xaxis: {
          title: { text: x },
          gridcolor: theme.gridColor,
          tickfont: { color: theme.fontColor },
        },
        yaxis: {
          title: { text: y },
          gridcolor: theme.gridColor,
          tickfont: { color: theme.fontColor },
        },
        zaxis: {
          title: { text: z },
          gridcolor: theme.gridColor,
          tickfont: { color: theme.fontColor },
        },
      };
      if (axisSizes) {
        sceneLayout.xaxis.range = [0, axisSizes.x];
        sceneLayout.yaxis.range = [0, axisSizes.y];
        sceneLayout.zaxis.range = [0, axisSizes.z];
      }

      Plotly.newPlot(
        plotRef.current!,
        [trace],
        {
          paper_bgcolor: theme.paperBg,
          plot_bgcolor: theme.plotBg,
          font: { color: theme.fontColor },
          scene: sceneLayout,
          margin: { l: 0, r: 0, t: 0, b: 0 },
        },
        { responsive: true }
      );
    })();

    return () => {
      cancelled = true;
    };
  }, [rows, x, y, z, plotType, colorscale, themeScheme]);


  function onFile(file: File) {
    papa.parse(file, {
      header: true,
      dynamicTyping: true,
      skipEmptyLines: true,
      complete: (res) => {
        const data = (res.data as Row[]).filter(
          (row) => row && Object.keys(row).length > 0
        );
        if (!data.length) {
          setRows([]);
          setColumns([]);
          setX("");
          setY("");
          setZ("");
          return;
        }
        const cols = Object.keys(data[0]);
        setRows(data);
        setColumns(cols);
        if (cols.length >= 3) {
          setX(cols[0]);
          setY(cols[1]);
          setZ(cols[2]);
        } else {
          setX("");
          setY("");
          setZ("");
        }
      },
    });
  }

  async function exportPNG() {
    const Plotly = await import("plotly.js-dist-min");

    if (!plotRef.current) return;

    Plotly.toImage(plotRef.current, {
      format: "png",
      width: 1600,
      height: 900,
    }).then((url) => {
      const a = document.createElement("a");
      a.href = url;
      a.download = "3dplot-gen.png";
      a.click();
    });
  }

  return (
    <div className="h-screen grid grid-cols-[280px_1fr] overflow-y-hidden">
      {/* Sidebar */}
      <aside className="bg-panel p-4 space-y-4 overflow-y-scroll">
        <div className="text-lg font-semibold">
          3dplot-gen
          <div className="text-xs text-muted">CSV → 3D Plot</div>
        </div>

        <Control label="CSV file">
          <input
            type="file"
            accept=".csv"
            onChange={(e) =>
              e.target.files && onFile(e.target.files[0])
            }
            className="w-full text-sm file:bg-neutral-900 file:text-slate-300 file:border-0
                       file:rounded file:px-2 file:py-1"
          />
        </Control>
        <p className="text-xs text-muted">
          First data row = axis sizes (X max, Y max, Z max). Plot uses rows after that.{" "}
          <a href="/example-surface.csv" className="text-accent underline" download>
            Example (surface)
          </a>
        </p>

        <Control label="X axis">
          <select
            className="w-full bg-neutral-900 border border-slate-700 rounded px-2 py-1"
            value={x}
            onChange={(e) => setX(e.target.value)}
          >
            <option value="">—</option>
            {columns.map((c) => (
              <option key={c} value={c}>
                {c}
              </option>
            ))}
          </select>
        </Control>

        <Control label="Y axis">
          <select
            className="w-full bg-neutral-900 border border-slate-700 rounded px-2 py-1"
            value={y}
            onChange={(e) => setY(e.target.value)}
          >
            <option value="">—</option>
            {columns.map((c) => (
              <option key={c} value={c}>
                {c}
              </option>
            ))}
          </select>
        </Control>

        <Control label="Z axis">
          <select
            className="w-full bg-neutral-900 border border-slate-700 rounded px-2 py-1"
            value={z}
            onChange={(e) => setZ(e.target.value)}
          >
            <option value="">—</option>
            {columns.map((c) => (
              <option key={c} value={c}>
                {c}
              </option>
            ))}
          </select>
        </Control>

        <Control label="Plot type">
          <select
            className="w-full bg-neutral-900 border border-slate-700 rounded px-2 py-1"
            value={plotType}
            onChange={(e) => setPlotType(e.target.value as PlotType)}
          >
            <option value="scatter">Dots (scatter)</option>
            <option value="surface">Surface (area)</option>
          </select>
        </Control>

        <Control label="Colorscale">
          <select
            className="w-full bg-neutral-900 border border-slate-700 rounded px-2 py-1"
            value={colorscale}
            onChange={(e) => setColorscale(e.target.value)}
          >
            {COLORSCALES.map((name) => (
              <option key={name} value={name}>
                {name}
              </option>
            ))}
          </select>
        </Control>

        <Control label="Background">
          <select
            className="w-full bg-neutral-900 border border-slate-700 rounded px-2 py-1"
            value={themeScheme}
            onChange={(e) => setThemeScheme(e.target.value as ThemeScheme)}
          >
            <option value="dark">Dark</option>
            <option value="light">Light</option>
            <option value="transparent">Transparent</option>
          </select>
        </Control>

        {plotType === "surface" &&
          plotData.length > 0 &&
          buildSurfaceGrid() === null && (
            <p className="text-xs text-amber-500/90">
              Surface needs a full X and Y grid of points. Showing dots instead.
            </p>
          )}

        <button
          onClick={exportPNG}
          className="w-full mt-2 bg-accent text-sm rounded px-2 py-1 hover:opacity-90"
        >
          Export PNG
        </button>
      </aside>

      {/* Plot */}
      <main className={THEME_SCHEMES[themeScheme].mainClass}>
        <div ref={plotRef} className="w-full h-full" />
      </main>
    </div>
  );
}

function Control({
  label,
  children,
}: {
  label: string;
  children: React.ReactNode;
}) {
  return (
    <div className="space-y-1">
      <label className="text-xs text-muted">{label}</label>
      {children}
    </div>
  );
}