breakpilot-lehrer/admin-lehrer/components/ocr-overlay/usePixelWordPositions.ts

import { useEffect, useState } from 'react'
import type { GridCell } from '@/app/(admin)/ai/ocr-overlay/types'

export interface WordPosition {
  xPct: number
  wPct: number
  text: string
  fontRatio: number
}

/**
 * Analyse dark-pixel clusters on an image to determine
 * the exact horizontal position & auto-font-size of word groups in each cell.
 *
 * When rotation=180, the image is rotated 180° before pixel analysis.
 * Cell coordinates are transformed to the rotated space for reading,
 * and cluster positions are mirrored back to the original coordinate system.
 *
 * Returns a Map<cell_id, WordPosition[]>.
 */
export function usePixelWordPositions(
  imageUrl: string,
  cells: GridCell[],
  active: boolean,
  rotation: 0 | 180 = 0,
): Map<string, WordPosition[]> {
  const [cellWordPositions, setCellWordPositions] = useState<Map<string, WordPosition[]>>(new Map())

  useEffect(() => {
    if (!active || cells.length === 0 || !imageUrl) return

    const img = new Image()
    img.crossOrigin = 'anonymous'
    img.onload = () => {
      const imgW = img.naturalWidth
      const imgH = img.naturalHeight

      const canvas = document.createElement('canvas')
      canvas.width = imgW
      canvas.height = imgH
      const ctx = canvas.getContext('2d')
      if (!ctx) return

      if (rotation === 180) {
        ctx.translate(imgW, imgH)
        ctx.rotate(Math.PI)
        ctx.drawImage(img, 0, 0)
        ctx.setTransform(1, 0, 0, 1, 0, 0)
      } else {
        ctx.drawImage(img, 0, 0)
      }

      const refFontSize = 40
      const fontFam = "'Liberation Sans', Arial, sans-serif"
      ctx.font = `${refFontSize}px ${fontFam}`

      const positions = new Map<string, WordPosition[]>()

      for (const cell of cells) {
        if (!cell.bbox_pct || !cell.text) continue

        const rawGroups = cell.text.split(/\s{3,}/).map(s => s.trim()).filter(Boolean)

        // Merge single-char symbol groups (OCR artifacts from box borders like "|", ">")
        // with their neighbour to avoid polluting the cluster-to-group matching
        const groups: string[] = []
        for (let gi = 0; gi < rawGroups.length; gi++) {
          const g = rawGroups[gi]
          const isArtifact = g.length <= 2 && !/[a-zA-Z0-9\u00C0-\u024F]/.test(g)
          if (isArtifact) {
            if (gi + 1 < rawGroups.length) {
              // merge with next group
              rawGroups[gi + 1] = g + '  ' + rawGroups[gi + 1]
            } else if (groups.length > 0) {
              // last group — merge with previous
              groups[groups.length - 1] += '  ' + g
            } else {
              groups.push(g)
            }
          } else {
            groups.push(g)
          }
        }

        let cx: number, cy: number
        const cw = Math.round(cell.bbox_pct.w / 100 * imgW)
        const ch = Math.round(cell.bbox_pct.h / 100 * imgH)

        if (rotation === 180) {
          cx = Math.round((100 - cell.bbox_pct.x - cell.bbox_pct.w) / 100 * imgW)
          cy = Math.round((100 - cell.bbox_pct.y - cell.bbox_pct.h) / 100 * imgH)
        } else {
          cx = Math.round(cell.bbox_pct.x / 100 * imgW)
          cy = Math.round(cell.bbox_pct.y / 100 * imgH)
        }
        if (cw <= 0 || ch <= 0) continue
        if (cx < 0) cx = 0
        if (cy < 0) cy = 0
        if (cx + cw > imgW || cy + ch > imgH) continue

        const imageData = ctx.getImageData(cx, cy, cw, ch)

        const proj = new Float32Array(cw)
        for (let y = 0; y < ch; y++) {
          for (let x = 0; x < cw; x++) {
            const idx = (y * cw + x) * 4
            const lum = 0.299 * imageData.data[idx] + 0.587 * imageData.data[idx + 1] + 0.114 * imageData.data[idx + 2]
            if (lum < 128) proj[x]++
          }
        }

        const threshold = Math.max(1, ch * 0.03)
        const minGap = Math.max(5, Math.round(cw * 0.02))
        let clusters: { start: number; end: number }[] = []
        let inCluster = false
        let clStart = 0
        let gap = 0

        for (let x = 0; x < cw; x++) {
          if (proj[x] >= threshold) {
            if (!inCluster) { clStart = x; inCluster = true }
            gap = 0
          } else if (inCluster) {
            gap++
            if (gap > minGap) {
              clusters.push({ start: clStart, end: x - gap })
              inCluster = false
              gap = 0
            }
          }
        }
        if (inCluster) clusters.push({ start: clStart, end: cw - 1 - gap })

        if (clusters.length === 0) continue

        // Filter out very narrow clusters (likely box borders / vertical lines)
        const minClusterW = Math.max(3, Math.round(cw * 0.005))
        clusters = clusters.filter(c => (c.end - c.start + 1) > minClusterW)
        if (clusters.length === 0) continue

        if (rotation === 180) {
          clusters = clusters.map(c => ({
            start: cw - 1 - c.end,
            end: cw - 1 - c.start,
          })).reverse()
        }

        const wordPos: WordPosition[] = []

        // Match groups to clusters using width-proportional assignment.
        // Each group is assigned to the cluster whose width best matches
        // the group's expected pixel width (text measurement).
        if (groups.length > 1 && clusters.length >= groups.length) {
          // Measure each group's expected width
          const groupWidths = groups.map(g => ctx.measureText(g).width)

          // Greedy assignment: for each group (in order), find the best
          // unassigned cluster by width ratio consistency
          const totalMeasured = groupWidths.reduce((a, b) => a + b, 0)
          const totalClusterW = clusters.reduce((a, c) => a + (c.end - c.start + 1), 0)
          const refScale = totalClusterW / totalMeasured
          const used = new Set<number>()

          const assignments: number[] = []
          for (let gi = 0; gi < groups.length; gi++) {
            const expectedW = groupWidths[gi] * refScale
            let bestIdx = -1
            let bestDiff = Infinity
            for (let ci = 0; ci < clusters.length; ci++) {
              if (used.has(ci)) continue
              const clW = clusters[ci].end - clusters[ci].start + 1
              const diff = Math.abs(clW - expectedW)
              if (diff < bestDiff) {
                bestDiff = diff
                bestIdx = ci
              }
            }
            used.add(bestIdx)
            assignments.push(bestIdx)
          }

          // Sort assignments to maintain left-to-right order
          const sortedPairs = assignments
            .map((ci, gi) => ({ ci, gi }))
            .sort((a, b) => clusters[a.ci].start - clusters[b.ci].start)

          for (const { ci, gi } of sortedPairs) {
            const cl = clusters[ci]
            const clusterW = cl.end - cl.start + 1
            const autoFontPx = refFontSize * (clusterW / groupWidths[gi])
            const fontRatio = Math.min(autoFontPx / ch, 1.0)
            wordPos.push({
              xPct: cell.bbox_pct.x + (cl.start / cw) * cell.bbox_pct.w,
              wPct: ((cl.end - cl.start + 1) / cw) * cell.bbox_pct.w,
              text: groups[gi],
              fontRatio,
            })
          }
        } else {
          // Single group OR not enough clusters:
          // use the WIDEST cluster (not first-to-last span which pulls in
          // stray pixels from adjacent page areas like box borders)
          const widest = clusters.reduce((best, c) =>
            (c.end - c.start) > (best.end - best.start) ? c : best, clusters[0])
          const clusterW = widest.end - widest.start + 1
          const measured = ctx.measureText(cell.text.trim())
          const autoFontPx = refFontSize * (clusterW / measured.width)
          const fontRatio = Math.min(autoFontPx / ch, 1.0)
          wordPos.push({
            xPct: cell.bbox_pct.x + (widest.start / cw) * cell.bbox_pct.w,
            wPct: ((widest.end - widest.start + 1) / cw) * cell.bbox_pct.w,
            text: cell.text.trim(),
            fontRatio,
          })
        }

        positions.set(cell.cell_id, wordPos)
      }

      // Normalise: find the most common fontRatio (mode) and apply it to all
      const allRatios: number[] = []
      for (const wps of positions.values()) {
        for (const wp of wps) allRatios.push(wp.fontRatio)
      }
      if (allRatios.length > 0) {
        const buckets = new Map<number, number>()
        for (const r of allRatios) {
          const key = Math.round(r * 50) / 50
          buckets.set(key, (buckets.get(key) || 0) + 1)
        }
        let modeRatio = allRatios[0]
        let modeCount = 0
        for (const [ratio, count] of buckets) {
          if (count > modeCount) { modeRatio = ratio; modeCount = count }
        }
        for (const wps of positions.values()) {
          for (const wp of wps) wp.fontRatio = modeRatio
        }
      }

      setCellWordPositions(positions)
    }
    img.src = imageUrl
  }, [active, cells, imageUrl, rotation])

  return cellWordPositions
}