import { useEffect, useMemo, useState } from "react";

export type PoolingKernel =
  | "Max (per channel)"
  | "Average (per channel)"
  | "Max (across channels)"
  | "Average (across channels)";

const MAX_SIZE = 500;

function imageDataToUrl(imageData: ImageData): string {
  const canvas = document.createElement("canvas");
  canvas.width = imageData.width;
  canvas.height = imageData.height;
  const context = canvas.getContext("2d");
  if (!context) throw new Error("Failed to get canvas context");
  context.putImageData(imageData, 0, 0);
  return canvas.toDataURL("image/png");
}

function grayscale(imageData: ImageData): ImageData {
  const output = new ImageData(new Uint8ClampedArray(imageData.data), imageData.width, imageData.height);
  for (let i = 0; i < output.data.length; i += 4) {
    const value = 0.299 * output.data[i] + 0.587 * output.data[i + 1] + 0.114 * output.data[i + 2];
    output.data[i] = output.data[i + 1] = output.data[i + 2] = value;
  }
  return output;
}

function pool(
  image: ImageData,
  kernelWidth: number,
  kernelHeight: number,
  stride: number,
  dilation: number,
  padding: number,
  kernel: PoolingKernel,
): ImageData {
  const effectiveWidth = dilation * (kernelWidth - 1) + 1;
  const effectiveHeight = dilation * (kernelHeight - 1) + 1;
  const width = Math.max(1, Math.floor((image.width + 2 * padding - effectiveWidth) / stride) + 1);
  const height = Math.max(1, Math.floor((image.height + 2 * padding - effectiveHeight) / stride) + 1);
  const output = new ImageData(width, height);
  const isMax = kernel.startsWith("Max");
  const acrossChannels = kernel.endsWith("(across channels)");

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const outputChannels = acrossChannels ? 1 : 3;
      for (let outputChannel = 0; outputChannel < outputChannels; outputChannel += 1) {
        let result = isMax ? -Infinity : 0;
        let count = 0;
        for (let ky = 0; ky < kernelHeight; ky += 1) {
          for (let kx = 0; kx < kernelWidth; kx += 1) {
            const sourceX = x * stride - padding + kx * dilation;
            const sourceY = y * stride - padding + ky * dilation;
            const firstChannel = acrossChannels ? 0 : outputChannel;
            const lastChannel = acrossChannels ? 3 : outputChannel + 1;
            for (let sourceChannel = firstChannel; sourceChannel < lastChannel; sourceChannel += 1) {
              const value = sourceX < 0 || sourceX >= image.width || sourceY < 0 || sourceY >= image.height
                ? 0
                : image.data[(sourceY * image.width + sourceX) * 4 + sourceChannel];
              result = isMax ? Math.max(result, value) : result + value;
              count += 1;
            }
          }
        }
        const value = isMax ? result : result / count;
        if (acrossChannels) {
          output.data[(y * width + x) * 4] = value;
          output.data[(y * width + x) * 4 + 1] = value;
          output.data[(y * width + x) * 4 + 2] = value;
        } else {
          output.data[(y * width + x) * 4 + outputChannel] = value;
        }
      }
      output.data[(y * width + x) * 4 + 3] = 255;
    }
  }
  return output;
}

export default function usePoolingProcessing(
  imageUrl: string,
  useColor: boolean,
  kernel: PoolingKernel,
  kernelWidth: number,
  kernelHeight: number,
  stride: number,
  dilation: number,
  padding: number,
) {
  const [source, setSource] = useState<ImageData | null>(null);

  useEffect(() => {
    let cancelled = false;
    const image = new Image();
    image.crossOrigin = "anonymous";
    image.onload = () => {
      const scale = Math.min(1, MAX_SIZE / image.naturalWidth, MAX_SIZE / image.naturalHeight);
      const canvas = document.createElement("canvas");
      canvas.width = Math.max(1, Math.round(image.naturalWidth * scale));
      canvas.height = Math.max(1, Math.round(image.naturalHeight * scale));
      const context = canvas.getContext("2d");
      if (!context) return;
      context.drawImage(image, 0, 0, canvas.width, canvas.height);
      if (!cancelled) setSource(context.getImageData(0, 0, canvas.width, canvas.height));
    };
    image.src = imageUrl;
    return () => { cancelled = true; };
  }, [imageUrl]);

  return useMemo(() => {
    if (!source) return { input: null, output: null, inputSize: null, outputSize: null };
    const input = useColor ? source : grayscale(source);
    const output = pool(input, kernelWidth, kernelHeight, stride, dilation, padding, kernel);
    return {
      input: imageDataToUrl(input),
      output: imageDataToUrl(output),
      inputSize: `${input.width} × ${input.height}`,
      outputSize: `${output.width} × ${output.height}`,
    };
  }, [source, useColor, kernel, kernelWidth, kernelHeight, stride, dilation, padding]);
}