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

const MAX_SIZE: number = 500;

function getBoundedDimensions(width: number, height: number): [number, number] {
  if (width <= MAX_SIZE && height <= MAX_SIZE) {
    return [width, height];
  }

  const scale = Math.min(MAX_SIZE / width, MAX_SIZE / height);
  const boundedWidth = Math.max(1, Math.round(width * scale));
  const boundedHeight = Math.max(1, Math.round(height * scale));
  return [boundedWidth, boundedHeight];
}

async function getImageData(imageUrl: string): Promise<ImageData> {
  const image = await new Promise<HTMLImageElement>((resolve, reject) => {
    const img = new Image();
    img.crossOrigin = "anonymous";
    img.onload = () => resolve(img);
    img.onerror = reject;
    img.src = imageUrl;
  })

  const sourceWidth = image.naturalWidth || image.width;
  const sourceHeight = image.naturalHeight || image.height;
  const [targetWidth, targetHeight] = getBoundedDimensions(sourceWidth, sourceHeight);

  const canvas = document.createElement("canvas")
  canvas.width = targetWidth;
  canvas.height = targetHeight;

  const ctx = canvas.getContext("2d");
  if (!ctx) {
    throw new Error("Failed to get canvas context");
  }

  ctx.drawImage(image, 0, 0, targetWidth, targetHeight);
  return ctx.getImageData(0, 0, canvas.width, canvas.height);
}


function getImageUrl(imageData: ImageData): string {
  const canvas = document.createElement("canvas");
  canvas.width = imageData.width;
  canvas.height = imageData.height;

  const ctx = canvas.getContext("2d")
  if (!ctx) {
    throw new Error("Failed to get canvas context");
  }

  ctx.putImageData(imageData, 0, 0);
  return canvas.toDataURL("image/png");
}

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

export type ConvolutionRange = {
  min: number;
  max: number;
};

type ConvolutionResult = ConvolutionRange & {
  values: Float32Array;
  width: number;
  height: number;
  channels: 1 | 3;
};

function getShape(array: number[][] | number[][][]): number[] {
  const shape: number[] = [];
  let current: unknown = array;

  while (Array.isArray(current)) {
    shape.push(current.length);
    current = current[0];
  }

  return shape;
}
function convolve(
  imageData: ImageData,
  kernel: number[][] | number[][][],
  useColor: boolean,
  stride: number,
  dilation: number,
  padding: number,
): ConvolutionResult {
	if (Array.isArray(kernel[0][0])) { // 3D kernel
		console.log('3D kernel', getShape(kernel))
		return convolveColor(imageData, kernel as number[][][], stride, dilation, padding);
	} else { // 2D kernel
		console.log('2D kernel', getShape(kernel))
		if (useColor) {
			return convolveColorDepthwise(imageData, kernel as number[][], stride, dilation, padding);
		} {
			return convolveGray(imageData, kernel as number[][], stride, dilation, padding);
		}
	}
}

function convolveGray(image: ImageData, kernel: number[][], stride: number, dilation: number, padding: number): ConvolutionResult {
  const kernelWidth = kernel[0].length;
  const kernelHeight = kernel.length;

  const width = image.width;
  const height = image.height;
  const inputData = image.data;

  const outputWidth = Math.max(1, Math.floor((width + 2 * padding - dilation * (kernelWidth - 1) - 1) / stride) + 1);
  const outputHeight = Math.max(1, Math.floor((height + 2 * padding - dilation * (kernelHeight - 1) - 1) / stride) + 1);
  const values = new Float32Array(outputWidth * outputHeight);
  let min = Infinity;
  let max = -Infinity;

  for (let y = 0; y < outputHeight; ++y) {
    for (let x = 0; x < outputWidth; ++x) {
      // dot product
      let sum = 0;
      for (let ky = 0; ky < kernelHeight; ++ky) {
        for (let kx = 0; kx < kernelWidth; ++kx) {
          const sourceX = x * stride - padding + kx * dilation;
          const sourceY = y * stride - padding + ky * dilation;
          if (sourceX < 0 || sourceX >= width || sourceY < 0 || sourceY >= height) continue;
          const pixelIndex = (sourceY * width + sourceX) * 4;
          const pixelValue = inputData[pixelIndex];
          const kernelValue = kernel[ky][kx];
          sum += pixelValue * kernelValue;
        }
      }

      values[y * outputWidth + x] = sum;
      min = Math.min(min, sum);
      max = Math.max(max, sum);
    }
  }

  return { values, width: outputWidth, height: outputHeight, channels: 1, min, max };
}

function convolveColor(image: ImageData, kernel: number[][][], stride: number, dilation: number, padding: number): ConvolutionResult {
  const kernelWidth = kernel[0][0].length;
  const kernelHeight = kernel[0].length;

  const width = image.width;
  const height = image.height;
  const inputData = image.data;

  const outputWidth = Math.max(1, Math.floor((width + 2 * padding - dilation * (kernelWidth - 1) - 1) / stride) + 1);
  const outputHeight = Math.max(1, Math.floor((height + 2 * padding - dilation * (kernelHeight - 1) - 1) / stride) + 1);
  const values = new Float32Array(outputWidth * outputHeight);
  let min = Infinity;
  let max = -Infinity;

  for (let y = 0; y < outputHeight; ++y) {
    for (let x = 0; x < outputWidth; ++x) {
      // dot product over 3 channels
      let sum = 0;
      for (let ky = 0; ky < kernelHeight; ++ky) {
        for (let kx = 0; kx < kernelWidth; ++kx) {
          const sourceX = x * stride - padding + kx * dilation;
          const sourceY = y * stride - padding + ky * dilation;
          if (sourceX < 0 || sourceX >= width || sourceY < 0 || sourceY >= height) continue;
          const pixelIndex = (sourceY * width + sourceX) * 4;
          const r = inputData[pixelIndex];
          const g = inputData[pixelIndex + 1];
          const b = inputData[pixelIndex + 2];

          const kernelR = kernel[0][ky][kx];
          const kernelG = kernel[1][ky][kx];
          const kernelB = kernel[2][ky][kx];
          sum += r * kernelR + g * kernelG + b * kernelB;
        }
      }

      values[y * outputWidth + x] = sum;
      min = Math.min(min, sum);
      max = Math.max(max, sum);
    }
  }

  return { values, width: outputWidth, height: outputHeight, channels: 1, min, max };
}

function convolveColorDepthwise(image: ImageData, kernel: number[][], stride: number, dilation: number, padding: number): ConvolutionResult {
  const kernelWidth = kernel[0].length;
  const kernelHeight = kernel.length;
  const outputWidth = Math.max(1, Math.floor((image.width + 2 * padding - dilation * (kernelWidth - 1) - 1) / stride) + 1);
  const outputHeight = Math.max(1, Math.floor((image.height + 2 * padding - dilation * (kernelHeight - 1) - 1) / stride) + 1);
  const values = new Float32Array(outputWidth * outputHeight * 3);
  let min = Infinity;
  let max = -Infinity;

	console.log('here');
  for (let y = 0; y < outputHeight; ++y) {
    for (let x = 0; x < outputWidth; ++x) {
      for (let channel = 0; channel < 3; ++channel) {
        let sum = 0;
        for (let ky = 0; ky < kernelHeight; ++ky) {
          for (let kx = 0; kx < kernelWidth; ++kx) {
            const sourceX = x * stride - padding + kx * dilation;
            const sourceY = y * stride - padding + ky * dilation;
            if (sourceX < 0 || sourceX >= image.width || sourceY < 0 || sourceY >= image.height) continue;
            const pixelIndex = (sourceY * image.width + sourceX) * 4 + channel;
            sum += image.data[pixelIndex] * kernel[ky][kx];
          }
        }
        values[(y * outputWidth + x) * 3 + channel] = sum;
        min = Math.min(min, sum);
        max = Math.max(max, sum);
      }
    }
  }

	console.log('done', outputHeight, outputWidth);

  return { values, width: outputWidth, height: outputHeight, channels: 3, min, max };
}

function convolutionToImageData(result: ConvolutionResult): ImageData {
  const outputData = new Uint8ClampedArray(result.width * result.height * 4);
  const range = result.max - result.min;

  const pixelCount = result.width * result.height;
  for (let i = 0; i < pixelCount; ++i) {
    const outputIndex = i * 4;
    for (let channel = 0; channel < 3; ++channel) {
      const rawValue = (result.channels === 1)
        ? result.values[i]
        : result.values[i * 3 + channel];
      outputData[outputIndex + channel] = (range === 0)
        ? 0
        : 255 * (rawValue - result.min) / range;
    }
    outputData[outputIndex + 3] = 255;
  }
	console.log('conversion done');

  return new ImageData(outputData, result.width, result.height);
}

export default function useConvolutionProcessing(
  rawInputImage: string,
  kernel: number[][][] | number[][],
  useColor: boolean,
  stride: number,
  dilation: number,
  padding: number,
): [string | null, string | null, ConvolutionRange | null] {
  const [rawImageData, setRawImageData] = useState<ImageData | null>(null);

  // extract input image data (array)
  useEffect(() => {
    let cancelled = false;

    async function processImage() {
      const imageData = await getImageData(rawInputImage);
      if (!cancelled) {
        setRawImageData(imageData);
      }
    }

    processImage();

    return () => {
      cancelled = true;
    }
  }, [rawInputImage]);

  const processedImageData = useMemo(() => {
    if (!rawImageData) return null;

    return useColor ? rawImageData : convertToGrayscale(rawImageData);
  }, [rawImageData, useColor]);

  const convolutionResult = useMemo(() => {
    if (!processedImageData) return null;

    return convolve(processedImageData, kernel, useColor, stride, dilation, padding);
  }, [processedImageData, kernel, useColor, stride, dilation, padding]);

  const inputImage = useMemo(() => {
    if (!processedImageData) return null;

    return getImageUrl(processedImageData);
  }, [processedImageData]);

  const outputImage = useMemo(() => {
    if (!convolutionResult) return null;

    return getImageUrl(convolutionToImageData(convolutionResult));
  }, [convolutionResult]);

	console.log('url');

  const outputRange = convolutionResult
    ? { min: convolutionResult.min, max: convolutionResult.max }
    : null;

	console.log('range');
  return [inputImage, outputImage, outputRange];
}