import { RawImage } from '../types';

/**
 * Convert RGB data to grayscale array.
 */
function toGrayscale(
  data: Buffer,
  width: number,
  height: number,
  channels: number
): Uint8Array {
  const gray = new Uint8Array(width * height);
  for (let i = 0; i < width * height; i++) {
    const idx = i * channels;
    gray[i] = Math.round(
      data[idx] * 0.299 + data[idx + 1] * 0.587 + data[idx + 2] * 0.114
    );
  }
  return gray;
}

/**
 * Produce a Sobel edge map from raw RGB data.
 */
export function createEdgeMap({
  data,
  width,
  height,
  channels,
}: RawImage): Uint8Array {
  const gray = toGrayscale(data, width, height, channels);
  const edges = new Uint8Array(width * height);

  for (let y = 1; y < height - 1; y++) {
    for (let x = 1; x < width - 1; x++) {
      const idx = y * width + x;
      const gx =
        -gray[(y - 1) * width + (x - 1)] +
        gray[(y - 1) * width + (x + 1)] -
        2 * gray[idx - 1] +
        2 * gray[idx + 1] -
        gray[(y + 1) * width + (x - 1)] +
        gray[(y + 1) * width + (x + 1)];

      const gy =
        -gray[(y - 1) * width + (x - 1)] -
        2 * gray[(y - 1) * width + x] -
        gray[(y - 1) * width + (x + 1)] +
        gray[(y + 1) * width + (x - 1)] +
        2 * gray[(y + 1) * width + x] +
        gray[(y + 1) * width + (x + 1)];

      const magnitude = Math.sqrt(gx * gx + gy * gy);
      edges[idx] = magnitude > 40 ? 1 : 0;
    }
  }

  return edges;
}

/**
 * Morphological closing (dilate followed by erode).
 */
export function morphologyClose(
  binary: Uint8Array,
  width: number,
  height: number,
  kernelSize: number
): Uint8Array {
  const dilated = dilate(binary, width, height, kernelSize);
  return erode(dilated, width, height, kernelSize);
}

export function dilate(
  binary: Uint8Array,
  width: number,
  height: number,
  kernelSize: number
): Uint8Array {
  const result = new Uint8Array(width * height);
  const offset = Math.floor(kernelSize / 2);

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      let maxVal = 0;

      for (let ky = -offset; ky <= offset; ky++) {
        for (let kx = -offset; kx <= offset; kx++) {
          const ny = y + ky;
          const nx = x + kx;

          if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
            maxVal = Math.max(maxVal, binary[ny * width + nx]);
          }
        }
      }

      result[y * width + x] = maxVal;
    }
  }

  return result;
}

export function erode(
  binary: Uint8Array,
  width: number,
  height: number,
  kernelSize: number
): Uint8Array {
  const result = new Uint8Array(width * height);
  const offset = Math.floor(kernelSize / 2);

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      let minVal = 1;

      for (let ky = -offset; ky <= offset; ky++) {
        for (let kx = -offset; kx <= offset; kx++) {
          const ny = y + ky;
          const nx = x + kx;

          if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
            minVal = Math.min(minVal, binary[ny * width + nx]);
          }
        }
      }

      result[y * width + x] = minVal;
    }
  }

  return result;
}

export { toGrayscale };