Mozaic/frontend/src/utils/floodFill.ts

import type { BoundingBox, Point, RegionMask, RGBColor } from '@/types/mosaic'
import { colorDistance, relativeLuminance } from './colors'
import { createRegionMask } from './regionMask'

export interface FloodFillOptions {
  tolerance: number
  darkBoundaryThreshold: number
  connectivity: 4 | 8
  occupancy?: Uint8Array
}

export interface FloodFillResult {
  mask: RegionMask
  pixelCount: number
  boundingBox: BoundingBox
  labelPosition: Point
  color: RGBColor
}

const pixelColor = (data: Uint8ClampedArray, index: number): RGBColor => ({
  r: data[index * 4] ?? 0,
  g: data[index * 4 + 1] ?? 0,
  b: data[index * 4 + 2] ?? 0,
})

const median = (values: number[]): number => {
  values.sort((first, second) => first - second)
  return values[Math.floor(values.length / 2)] ?? 0
}

export function floodFillRegion(
  image: ImageData,
  startX: number,
  startY: number,
  options: FloodFillOptions,
): FloodFillResult | null {
  const { width, height, data } = image
  if (startX < 0 || startY < 0 || startX >= width || startY >= height) return null
  const startIndex = startY * width + startX
  if (options.occupancy?.[startIndex]) return null

  const seed = pixelColor(data, startIndex)
  const seedIsDark = relativeLuminance(seed) <= options.darkBoundaryThreshold
  const visited = new Uint8Array(width * height)
  const stack: number[] = [startIndex]
  const runs: number[] = []
  const reds: number[] = []
  const greens: number[] = []
  const blues: number[] = []
  let pixelCount = 0
  let minX = width
  let minY = height
  let maxX = 0
  let maxY = 0

  const isAccepted = (index: number) => {
    if (visited[index] || options.occupancy?.[index]) return false
    const color = pixelColor(data, index)
    if (!seedIsDark && relativeLuminance(color) <= options.darkBoundaryThreshold) return false
    return colorDistance(color, seed) <= options.tolerance
  }

  const sample = (index: number) => {
    const color = pixelColor(data, index)
    if (reds.length < 8000) {
      reds.push(color.r)
      greens.push(color.g)
      blues.push(color.b)
      return
    }
    const replacement = Math.floor(Math.random() * pixelCount)
    if (replacement < 8000) {
      reds[replacement] = color.r
      greens[replacement] = color.g
      blues[replacement] = color.b
    }
  }

  while (stack.length) {
    const index = stack.pop()
    if (index === undefined || !isAccepted(index)) continue
    const y = Math.floor(index / width)
    let left = index % width
    let right = left
    while (left > 0 && isAccepted(y * width + left - 1)) left -= 1
    while (right + 1 < width && isAccepted(y * width + right + 1)) right += 1

    for (let x = left; x <= right; x += 1) {
      const pixelIndex = y * width + x
      visited[pixelIndex] = 1
      pixelCount += 1
      // Avoid both ends of the scanline, where outlines and JPEG artifacts are most likely.
      if (x >= left + 2 && x <= right - 2) sample(pixelIndex)
    }
    runs.push(y, left, right)
    minX = Math.min(minX, left)
    maxX = Math.max(maxX, right)
    minY = Math.min(minY, y)
    maxY = Math.max(maxY, y)

    for (const neighborY of [y - 1, y + 1]) {
      if (neighborY < 0 || neighborY >= height) continue
      const extra = options.connectivity === 8 ? 1 : 0
      const scanStart = Math.max(0, left - extra)
      const scanEnd = Math.min(width - 1, right + extra)
      let insideCandidate = false
      for (let x = scanStart; x <= scanEnd; x += 1) {
        const neighborIndex = neighborY * width + x
        const candidate = isAccepted(neighborIndex)
        if (candidate && !insideCandidate) stack.push(neighborIndex)
        insideCandidate = candidate
      }
    }
  }
  if (!pixelCount) return null
  if (!reds.length) {
    reds.push(seed.r)
    greens.push(seed.g)
    blues.push(seed.b)
  }

  const orderedRuns: Array<[number, number, number]> = []
  for (let index = 0; index < runs.length; index += 3) {
    orderedRuns.push([runs[index] ?? 0, runs[index + 1] ?? 0, runs[index + 2] ?? 0])
  }
  orderedRuns.sort((first, second) => first[0] - second[0] || first[1] - second[1])
  const centerY = (minY + maxY) / 2
  const labelRun = orderedRuns.reduce((best, run) => {
    const score = run[2] - run[1] - Math.abs(run[0] - centerY) * 0.2
    const bestScore = best[2] - best[1] - Math.abs(best[0] - centerY) * 0.2
    return score > bestScore ? run : best
  }, orderedRuns[0] ?? [startY, startX, startX])

  return {
    mask: createRegionMask(width, height, orderedRuns.flatMap((run) => run)),
    pixelCount,
    boundingBox: { x: minX, y: minY, width: maxX - minX + 1, height: maxY - minY + 1 },
    labelPosition: { x: Math.round((labelRun[1] + labelRun[2]) / 2), y: labelRun[0] },
    color: { r: median(reds), g: median(greens), b: median(blues) },
  }
}