src/ConvolutionVisualizer.tsx

import { useEffect, useState, useRef } from "react";
import marioImage from "./assets/mario.png";
import useConvolutionProcessing from "./useConvolutionProcessing.ts";
import { Button, Card, Dropdown, Radio } from "@elvis/ui";

import { 
  DEFAULT_COLOR_KERNEL, 
  DEFAULT_GRAY_KERNEL,
  GRAY_KERNEL_PRESETS,
  COLOR_KERNEL_PRESETS,
} from "./kernels.ts";

const DEFAULT_IMAGE: string = marioImage;
const MIN_KERNEL_SIZE = 1;
const MAX_KERNEL_SIZE = 20;

export default function ConvolutionVisualizer() {
  const [rawInputImage, setRawInputImage] = useState<string>(DEFAULT_IMAGE);
  const [useColor, setUseColor] = useState<boolean>(true);
  const [colorKernel, setColorKernel] = useState<number[][][]>(DEFAULT_COLOR_KERNEL);
  const [grayscaleKernel, setGrayscaleKernel] = useState<number[][]>(DEFAULT_GRAY_KERNEL);
  const kernel = useColor ? colorKernel : grayscaleKernel;

  const [inputImage, outputImage] = useConvolutionProcessing(rawInputImage, kernel);

  return (
    <div className="grid grid-cols-2 gap-8 h-full min-h-0">
      <InputOutputViewer 
        input={inputImage}
        output={outputImage}
      />

      <div className="flex flex-col gap-8 h-full min-h-0">
        <ImageControls
          setImage={setRawInputImage}
          useColor={useColor}
          setUseColor={setUseColor}
        />
        <KernelEditor
          useColor={useColor}
          colorKernel={colorKernel}
          setColorKernel={setColorKernel}
          grayscaleKernel={grayscaleKernel}
          setGrayscaleKernel={setGrayscaleKernel}
        />
      </div>
    </div>
  );
}

type InputOutputViewerProps = {
  input: string | null;
  output: string | null;
}
function InputOutputViewer({ input, output }: InputOutputViewerProps) {
  return (
    <Card className="grid grid-rows-2 h-full min-h-0">

      <div className="flex flex-col items-center justify-center min-h-0">
        <h2 className="text-lg font-bold mb-2">Input Image</h2>
        {input ? (
          <img src={input} alt="Input" className="flex-1 max-w-full max-h-full min-h-0" />
        ) : (
          <p>Loading...</p>
        )}
      </div>

      <div className="flex flex-col items-center justify-center min-h-0">
        <h2 className="text-lg font-bold mb-2">Output Image</h2>
        {output ? (
          <img src={output} alt="Output" className="flex-1 max-w-full max-h-full min-h-0" />
        ) : (
          <p>Processing...</p>
        )}
      </div>

    </Card>
  );
}

type ImageControlsProps = {
  setImage: (imageUrl: string) => void;
  useColor: boolean;
  setUseColor: (useColor: boolean) => void;
}
function ImageControls({ setImage, useColor, setUseColor }: ImageControlsProps) {
  const imageFileInputRef = useRef<HTMLInputElement | null>(null);

  return (
    <Card className="flex justify-evenly items-center p-4 gap-4">
      <input
        ref={imageFileInputRef}
        type="file"
        accept="image/*"
        onChange={(e) => {
          const file = e.target.files?.[0];
          if (!file) return;

          const reader = new FileReader();
          reader.onload = () => {
            const result = reader.result;
            if (typeof result !== "string") {
              return;
            }
            setImage(result);
          };
          reader.readAsDataURL(file);
        }}
        style={{ display: "none" }}
      />
      <Button
        label="Upload Image"
        onClick={() => imageFileInputRef.current?.click()}
      />
      <Radio
        label="Color option"
        options={["Grayscale", "Color"] as const}
        activeOption={useColor ? "Color" : "Grayscale"}
        onChange={(option) => setUseColor(option === "Color")}
      />
    </Card>
  );
}

type KernelEditorProps = {
  useColor: boolean;
  colorKernel: number[][][];
  setColorKernel: (kernel: number[][][]) => void;
  grayscaleKernel: number[][];
  setGrayscaleKernel: (kernel: number[][]) => void;
}
function KernelEditor({
  useColor,
  colorKernel,
  setColorKernel,
  grayscaleKernel,
  setGrayscaleKernel,
}: KernelEditorProps) {
  const colorPresetNames = Object.keys(COLOR_KERNEL_PRESETS) as string[];
  const grayPresetNames = Object.keys(GRAY_KERNEL_PRESETS) as string[];
  const [selectedChannel, setSelectedChannel] = useState<number>(0);
  const [selectedColorPreset, setSelectedColorPreset] = useState<string>(colorPresetNames[0]);
  const [selectedGrayPreset, setSelectedGrayPreset] = useState<string>(grayPresetNames[0]);
  const [loadedColorPreset, setLoadedColorPreset] = useState<string | null>("Laplacian");
  const [loadedGrayPreset, setLoadedGrayPreset] = useState<string | null>("Laplacian");

  function colorKernelToRepr(nextKernel: number[][][]): string[][][] {
    return nextKernel.map((channel) =>
      channel.map((row) => row.map((value) => value.toString())),
    );
  }

  function grayKernelToRepr(nextKernel: number[][]): string[][] {
    return nextKernel.map((row) => row.map((value) => value.toString()));
  }

  const [draftColorKernel, setDraftColorKernel] = useState<string[][][]>(colorKernelToRepr(colorKernel));
  const [draftGrayKernel, setDraftGrayKernel] = useState<string[][]>(grayKernelToRepr(grayscaleKernel));

  useEffect(() => {
    setDraftColorKernel(colorKernelToRepr(colorKernel));
  }, [colorKernel]);

  useEffect(() => {
    setDraftGrayKernel(grayKernelToRepr(grayscaleKernel));
  }, [grayscaleKernel]);

  useEffect(() => {
    setSelectedChannel(0);
  }, [useColor]);

  function parseCell(cell: string): number | null {
    const trimmed = cell.trim();
    if (trimmed === "") return null;
    const parsed = Number(trimmed);
    return Number.isFinite(parsed) ? parsed : null;
  }

  function cloneColorKernel(nextKernel: number[][][]): number[][][] {
    return nextKernel.map((channel) => channel.map((row) => [...row]));
  }

  function cloneGrayKernel(nextKernel: number[][]): number[][] {
    return nextKernel.map((row) => [...row]);
  }

  function handleLoadPreset() {
    if (useColor) {
      const presetKernel = cloneColorKernel(
        COLOR_KERNEL_PRESETS[selectedColorPreset as keyof typeof COLOR_KERNEL_PRESETS],
      );
      setDraftColorKernel(colorKernelToRepr(presetKernel));
      setColorKernel(presetKernel);
      setLoadedColorPreset(selectedColorPreset);
      setSelectedChannel(0);
      return;
    }

    const presetKernel = cloneGrayKernel(
      GRAY_KERNEL_PRESETS[selectedGrayPreset as keyof typeof GRAY_KERNEL_PRESETS],
    );
    setDraftGrayKernel(grayKernelToRepr(presetKernel));
    setGrayscaleKernel(presetKernel);
    setLoadedGrayPreset(selectedGrayPreset);
  }

  function kernelsEqual2D(a: number[][], b: number[][]): boolean {
    if (a.length !== b.length) return false;
    if ((a[0]?.length ?? 0) !== (b[0]?.length ?? 0)) return false;
    return a.every((row, rowIndex) =>
      row.every((value, colIndex) => value === b[rowIndex][colIndex]),
    );
  }

  function kernelsEqual3D(a: number[][][], b: number[][][]): boolean {
    if (a.length !== b.length) return false;
    return a.every((channel, channelIndex) => kernelsEqual2D(channel, b[channelIndex]));
  }

  function handleKernelChange(channel: number, row: number, col: number, value: string) {
    if (useColor) {
      const nextDraft = draftColorKernel.map((c) => c.map((r) => [...r]));
      nextDraft[channel][row][col] = value;
      setDraftColorKernel(nextDraft);

      const parsedKernel = nextDraft.map((c) => c.map((r) => r.map(parseCell)));
      const isValid = parsedKernel.every((c) => c.every((r) => r.every((v) => v !== null)));
      if (isValid) {
        setColorKernel(parsedKernel as number[][][]);
      }
      return;
    }

    const nextDraft = draftGrayKernel.map((r) => [...r]);
    nextDraft[row][col] = value;
    setDraftGrayKernel(nextDraft);

    const parsedKernel = nextDraft.map((r) => r.map(parseCell));
    const isValid = parsedKernel.every((r) => r.every((v) => v !== null));
    if (isValid) {
      setGrayscaleKernel(parsedKernel as number[][]);
    }
  }

  function handleSizeChange(newWidth: number, newHeight: number) {
    const clampedWidth = Math.min(MAX_KERNEL_SIZE, Math.max(MIN_KERNEL_SIZE, newWidth));
    const clampedHeight = Math.min(MAX_KERNEL_SIZE, Math.max(MIN_KERNEL_SIZE, newHeight));

    function resizeMatrix(matrix: string[][], width: number, height: number): string[][] {
      return Array.from({ length: height }, (_, rowIndex) =>
        Array.from({ length: width }, (_, colIndex) => matrix[rowIndex]?.[colIndex] ?? "0"),
      );
    }

    if (useColor) {
      const resizedDraft = draftColorKernel.map((channelMatrix) =>
        resizeMatrix(channelMatrix, clampedWidth, clampedHeight),
      );
      setDraftColorKernel(resizedDraft);
      setColorKernel(
        resizedDraft.map((channelMatrix) =>
          channelMatrix.map((row) => row.map((cell) => parseCell(cell) ?? 0)),
        ),
      );
      return;
    }

    const resizedDraft = resizeMatrix(draftGrayKernel, clampedWidth, clampedHeight);
    setDraftGrayKernel(resizedDraft);
    setGrayscaleKernel(resizedDraft.map((row) => row.map((cell) => parseCell(cell) ?? 0)));
  }

  const channelLabels = ["R", "G", "B"];
  const activeMatrix = useColor
    ? draftColorKernel[selectedChannel]
    : draftGrayKernel;
  const currentHeight = activeMatrix.length;
  const currentWidth = activeMatrix[0]?.length ?? 0;
  const isModified = useColor
    ? loadedColorPreset !== null
      && !kernelsEqual3D(
        colorKernel,
        COLOR_KERNEL_PRESETS[loadedColorPreset as keyof typeof COLOR_KERNEL_PRESETS],
      )
    : loadedGrayPreset !== null
      && !kernelsEqual2D(
        grayscaleKernel,
        GRAY_KERNEL_PRESETS[loadedGrayPreset as keyof typeof GRAY_KERNEL_PRESETS],
      );
  const loadedPresetLabel = useColor ? loadedColorPreset : loadedGrayPreset;

  return (
    <Card className="flex flex-col items-center w-full min-h-0 h-full min-h-0 overflow-auto">
      <div className="flex items-end gap-2 mb-3">
        <Dropdown
          label="Kernel Preset"
          options={useColor ? colorPresetNames : grayPresetNames}
          activeOption={useColor ? selectedColorPreset : selectedGrayPreset}
          onChange={(option) => {
            if (useColor) {
              setSelectedColorPreset(option);
              return;
            }
            setSelectedGrayPreset(option);
          }}
        />
        <Button label="Load preset" onClick={handleLoadPreset} />
      </div>
      {loadedPresetLabel && (
        <div className="text-sm mb-3">
          <span>Loaded: {loadedPresetLabel}</span>
          {isModified && <span className="text-orange-700 ml-2">(modified)</span>}
        </div>
      )}
      { useColor && (
        <div className="flex gap-2 mb-3">
          {channelLabels.map((label, index) => (
            <button
              key={label}
              type="button"
              className={`px-3 py-1 text-sm rounded border ${
                selectedChannel === index ? "bg-orange-200 border-orange-300" : "bg-white border-gray-300"
              }`}
              onClick={() => setSelectedChannel(index)}
            >
              {label}
            </button>
          ))}
        </div>
      )}
      <div className="flex flex-col gap-4 mb-3">
        <div className="flex items-center gap-2">
          <span className="text-sm">Width</span>
          <button
            type="button"
            className="px-2 py-1 text-sm rounded border border-gray-300"
            onClick={() => handleSizeChange(currentWidth - 1, currentHeight)}
          >
            -
          </button>
          <span className="text-sm w-6 text-center">{currentWidth}</span>
          <button
            type="button"
            className="px-2 py-1 text-sm rounded border border-gray-300"
            onClick={() => handleSizeChange(currentWidth + 1, currentHeight)}
          >
            +
          </button>
        </div>
        <div className="flex items-center gap-2">
          <span className="text-sm">Height</span>
          <button
            type="button"
            className="px-2 py-1 text-sm rounded border border-gray-300"
            onClick={() => handleSizeChange(currentWidth, currentHeight - 1)}
          >
            -
          </button>
          <span className="text-sm w-6 text-center">{currentHeight}</span>
          <button
            type="button"
            className="px-2 py-1 text-sm rounded border border-gray-300"
            onClick={() => handleSizeChange(currentWidth, currentHeight + 1)}
          >
            +
          </button>
        </div>
      </div>
      <div className="flex flex-col gap-2">
        {activeMatrix.map((row, rowIndex) => (
          <div key={rowIndex} className="flex gap-2">
            {row.map((cellValue, colIndex) => (
              <input
                key={`${rowIndex}-${colIndex}`}
                type="text"
                className="w-14 px-2 py-1 border border-gray-300 rounded text-sm"
                value={cellValue}
                onChange={(event) =>
                  handleKernelChange(
                    useColor ? selectedChannel : 0,
                    rowIndex,
                    colIndex,
                    event.target.value,
                  )
                }
              />
            ))}
          </div>
        ))}
      </div>
    </Card>
  );
}