import { PaddingMode, ConvolutionStep, ImageDataExtended } from '@/types'

export function getPixel(
  image: ImageDataExtended,
  x: number,
  y: number,
  channel: number,
  padding: PaddingMode
): number {
  const { width, height, channels, data } = image

  let px = x
  let py = y

  switch (padding) {
    case 'zero':
      if (px < 0 || px >= width || py < 0 || py >= height) return 0
      break
    case 'replicate':
      px = Math.max(0, Math.min(px, width - 1))
      py = Math.max(0, Math.min(py, height - 1))
      break
    case 'reflect':
      if (px < 0) px = -px
      if (px >= width) px = 2 * width - px - 2
      if (py < 0) py = -py
      if (py >= height) py = 2 * height - py - 2
      px = Math.max(0, Math.min(px, width - 1))
      py = Math.max(0, Math.min(py, height - 1))
      break
    case 'wrap':
      px = ((px % width) + width) % width
      py = ((py % height) + height) % height
      break
  }

  const index = (py * width + px) * channels + channel
  return data[index] / 255
}

export function setPixel(
  image: ImageDataExtended,
  x: number,
  y: number,
  channel: number,
  value: number
): void {
  const { width, channels, data } = image
  if (x < 0 || x >= width || y < 0 || y >= image.height) return
  const index = (y * width + x) * channels + channel
  data[index] = Math.max(0, Math.min(255, Math.round(value * 255)))
}

export function* convolutionSteps(
  image: ImageDataExtended,
  kernel: number[][],
  padding: PaddingMode,
  normalize: boolean,
  onProgress?: (progress: number) => void
): Generator<ConvolutionStep> {
  const { width, height, channels } = image
  const kernelSize = kernel.length
  const halfKernel = Math.floor(kernelSize / 2)

  const totalPixels = width * height * channels
  let processedPixels = 0

  const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
  const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

  const outputData = new Float32Array(width * height * channels)

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      for (let c = 0; c < channels; c++) {
        const inputPatch: number[][] = []
        const multiplications: number[][] = []
        let sum = 0

        for (let ky = 0; ky < kernelSize; ky++) {
          const patchRow: number[] = []
          const multRow: number[] = []
          for (let kx = 0; kx < kernelSize; kx++) {
            const px = x + kx - halfKernel
            const py = y + ky - halfKernel
            const pixelValue = getPixel(image, px, py, c, padding)
            const kernelValue = kernel[ky][kx]
            patchRow.push(pixelValue)
            const mult = pixelValue * kernelValue
            multRow.push(mult)
            sum += mult
          }
          inputPatch.push(patchRow)
          multiplications.push(multRow)
        }

        const outputValue = sum
        const normalizedValue = sum * normalizeFactor

        const outputIndex = (y * width + x) * channels + c
        outputData[outputIndex] = normalizedValue

        processedPixels++
        const progress = processedPixels / totalPixels

        yield {
          x,
          y,
          channel: c,
          inputPatch,
          kernel,
          multiplications,
          sum: outputValue,
          outputValue: normalizedValue,
          normalizedValue,
          progress,
        }

        if (onProgress) onProgress(progress)
      }
    }
  }

  // Convert output to ImageData
  const outputImageData = new ImageData(width, height)
  for (let i = 0; i < outputData.length; i++) {
    outputImageData.data[i] = Math.max(0, Math.min(255, Math.round(outputData[i] * 255)))
  }

  // Final yield with complete output
  yield {
    x: -1,
    y: -1,
    channel: -1,
    inputPatch: [],
    kernel,
    multiplications: [],
    sum: 0,
    outputValue: 0,
    normalizedValue: 0,
    progress: 1,
  }
}

export function convolveImage(
  imageData: ImageData,
  kernel: number[][],
  padding: PaddingMode = 'zero',
  normalize: boolean = true
): ImageData {
  const extended: ImageDataExtended = {
    data: imageData.data,
    width: imageData.width,
    height: imageData.height,
    channels: 4,
  }

  const { width, height } = imageData
  const kernelSize = kernel.length
  const halfKernel = Math.floor(kernelSize / 2)
  const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
  const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

  const outputData = new Uint8ClampedArray(imageData.data.length)

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      for (let c = 0; c < 3; c++) {
        let sum = 0
        for (let ky = 0; ky < kernelSize; ky++) {
          for (let kx = 0; kx < kernelSize; kx++) {
            const px = x + kx - halfKernel
            const py = y + ky - halfKernel
            const pixelValue = getPixel(extended, px, py, c, padding)
            sum += pixelValue * kernel[ky][kx]
          }
        }
        const outputIndex = (y * width + x) * 4 + c
        outputData[outputIndex] = Math.max(0, Math.min(255, Math.round(sum * normalizeFactor * 255)))
      }
      outputData[(y * width + x) * 4 + 3] = 255
    }
  }

  return new ImageData(outputData, imageData.width, imageData.height)
}