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

/**
 * Computes a single convolution step for a specific pixel position.
 * Useful for static visualization/inspection without running the full animation.
 */
export function computeConvolutionStep(
  image: ImageDataExtended,
  kernel: number[][],
  padding: PaddingMode,
  normalize: boolean,
  x: number,
  y: number,
  channel: number
): ConvolutionStep {
  const { width, height, channels } = image
  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 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, channel, 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

  // Calculate progress (for consistency with generator)
  const totalPixels = width * height * channels
  const processedPixels = (y * width + x) * channels + channel + 1
  const progress = processedPixels / totalPixels

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

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 progressive rendering: yield partial output every N rows
  const rowsPerYield = Math.max(1, Math.floor(height / 20)) // ~20 yields per image
  let rowsSinceYield = 0

  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

        // Create partial output image for progressive rendering
        let partialOutputImageData: ImageData | undefined
        let isPartialOutput = false

        if (c === channels - 1 && x === width - 1) {
          // End of row - check if we should yield partial output
          rowsSinceYield++
          if (rowsSinceYield >= rowsPerYield || y === height - 1) {
            // Create partial output image
            const partialImageData = new ImageData(width, height)
            for (let i = 0; i < outputData.length; i++) {
              // Only fill up to current progress
              const pixelIdx = Math.floor(i / channels)
              const currentPixelIdx = y * width + x
              if (pixelIdx <= currentPixelIdx) {
                partialImageData.data[i] = Math.max(0, Math.min(255, Math.round(outputData[i] * 255)))
              } else {
                // Unprocessed pixels - show as transparent or original
                partialImageData.data[i] = 0
                partialImageData.data[i + 3] = 0 // transparent alpha for unprocessed
              }
            }
            // Fill alpha for processed pixels
            for (let py = 0; py <= y; py++) {
              for (let px = 0; px < width; px++) {
                for (let pc = 0; pc < channels; pc++) {
                  const idx = (py * width + px) * channels + pc
                  const imgIdx = (py * width + px) * 4 + pc
                  if (pc < 3) {
                    partialImageData.data[imgIdx] = Math.max(0, Math.min(255, Math.round(outputData[idx] * 255)))
                  } else {
                    partialImageData.data[imgIdx] = 255
                  }
                }
              }
            }

            partialOutputImageData = partialImageData
            isPartialOutput = true
            rowsSinceYield = 0
          }
        }

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

        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,
    partialOutputImageData: outputImageData,
    isPartialOutput: true,
  }
}

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)
}