ragflow/internal/deepdoc/parser/pdf/table/table_cells.go

package table

import (
	"log/slog"
	"math"
	pdf "ragflow/internal/deepdoc/parser/pdf/type"
	"ragflow/internal/deepdoc/parser/pdf/util"
	"regexp"
	"sort"
	"strings"
)

// ── TSR cell grouping ──────────────────────────────────────────────────

// GroupTSRCellsToRows groups TSR cells into rows by Y proximity.
// This is the basic fallback grouping used when model-specific grouping
// (e.g. EE label-aware grouping) is not applicable.
func GroupTSRCellsToRows(cells []pdf.TSRCell) [][]pdf.TSRCell {
	if len(cells) == 0 {
		return nil
	}
	if len(cells) == 1 {
		return [][]pdf.TSRCell{{cells[0]}}
	}
	heights := make([]float64, len(cells))
	for i, c := range cells {
		heights[i] = c.Y1 - c.Y0
	}
	sort.Float64s(heights)
	medianH := heights[len(heights)/2]
	if medianH <= 0 {
		medianH = 10
	}
	rowThreshold := medianH * 0.5

	sort.Slice(cells, func(i, j int) bool {
		if math.Abs(cells[i].Y0-cells[j].Y0) < rowThreshold {
			return cells[i].X0 < cells[j].X0
		}
		return cells[i].Y0 < cells[j].Y0
	})

	var rows [][]pdf.TSRCell
	var curRow []pdf.TSRCell
	curY := 0.0
	for _, c := range cells {
		if len(curRow) == 0 {
			curRow = append(curRow, c)
			curY = c.Y0
			continue
		}
		if c.Y0-curY > rowThreshold {
			rows = append(rows, curRow)
			curRow = []pdf.TSRCell{c}
			curY = c.Y0
		} else {
			curRow = append(curRow, c)
		}
	}
	if len(curRow) > 0 {
		rows = append(rows, curRow)
	}
	for _, row := range rows {
		sort.Slice(row, func(i, j int) bool { return row[i].X0 < row[j].X0 })
	}
	return rows
}

// ── cell text filling ──────────────────────────────────────────────────

// FillCellTextFromBoxes assigns PDF text boxes to TSR grid cells, mirroring
// Python's construct_table box→cell assignment (pdf_parser.py +
// table_structure_recognizer.py):
//
//  1. For each box, pick the single BEST row by vertical-overlap ratio
//     inter(box,rowStrip)/area(box) >= 0.3, tie-broken by inter/rowArea
//     (Python find_overlapped_with_threshold on the full-width row strip).
//  2. Within that row, pick the TIGHTEST column by horizontal edge/center
//     distance, requiring vertical overlap (Python find_horizontally_tightest_fit,
//     NO threshold). The box lands in exactly ONE cell (R,C).
//  3. Multiple boxes mapped to the same (R,C) are concatenated (Python joins
//     them in construct_table).
//
// This replaces the old many-to-many 2-D cell-overlap filter
// (inter(box,cell)/area(box) >= 0.3 on every cross-product cell), which
// duplicated a straddling box into two cells (#1), dropped boxes whose 2-D
// cell overlap was < 0.3 even though their row-vertical overlap was >= 0.3
// (#2), and never applied the inter/cellArea tie-break (#3). All three are
// go_bug in testdata/parity/known_diffs.json.
//
// The Go-only 0.85 guard (BoxMatchesCell) is retained for PRE-FILLED cells
// only: if a cell already carries text (e.g. per-cell OCR in the rotated
// path), a detected box overrides it only when it sits almost entirely inside
// the cell (>= 0.85). Empty cells accept any box the row/column selection
// picked, matching Python. See go_intentional rule
// table-cell-fill-filled-threshold-0.85.
func FillCellTextFromBoxes(cells []pdf.TSRCell, boxes []pdf.TextBox) {
	slog.Debug("fillCellTextFromBoxes", "cells", len(cells), "boxes", len(boxes))
	if len(cells) == 0 || len(boxes) == 0 {
		return
	}

	// Group cells into row bands by their top coordinate. The grid is a TSR
	// row×column cross-product, so every cell in a row shares the same Y band.
	// A row band spans the full table width (union of its cells), matching
	// Python's full-width "table row" components.
	type rowBand struct {
		y0, y1  float64
		stripX0 float64
		stripX1 float64
		cells   []int // indices into `cells`
	}
	var rows []rowBand
	// Group cells into the same row band by exact top coordinate. A TSR
	// cross-product grid (GroupCells) assigns every cell in a row the SAME
	// Y0 value, and distinct rows differ by at least a row height, so a tiny
	// epsilon is enough and never merges two real rows.
	const yTol = 1e-6
	for i := range cells {
		c := &cells[i]
		if c.X1 <= c.X0 || c.Y1 <= c.Y0 {
			continue // degenerate / span-covered cell: not a fill target
		}
		rb := &rowBand{}
		found := false
		for ri := range rows {
			if math.Abs(rows[ri].y0-c.Y0) <= yTol {
				rb = &rows[ri]
				found = true
				break
			}
		}
		if !found {
			rows = append(rows, rowBand{
				y0: c.Y0, y1: c.Y1, stripX0: c.X0, stripX1: c.X1,
			})
			rb = &rows[len(rows)-1]
		}
		if c.X0 < rb.stripX0 {
			rb.stripX0 = c.X0
		}
		if c.X1 > rb.stripX1 {
			rb.stripX1 = c.X1
		}
		if c.Y1 > rb.y1 {
			rb.y1 = c.Y1
		}
		rb.cells = append(rb.cells, i)
	}
	// Stable ordering: rows top-to-bottom, cells left-to-right (matches
	// Python's first-wins tie-breaking in find_overlapped_with_threshold /
	// find_horizontally_tightest_fit).
	sort.Slice(rows, func(i, j int) bool { return rows[i].y0 < rows[j].y0 })
	for ri := range rows {
		sort.Slice(rows[ri].cells, func(a, b int) bool {
			return cells[rows[ri].cells[a]].X0 < cells[rows[ri].cells[b]].X0
		})
	}

	// Accumulate box text per target cell so multiple boxes in one cell join.
	cellText := make([]string, len(cells))
	cellFilled := make([]bool, len(cells))
	matched := 0

	for bi := range boxes {
		b := boxes[bi]
		if IsCaptionBox(b.Text, b.LayoutType) {
			continue
		}
		boxArea := util.Area(&b)
		if boxArea <= 0 {
			continue
		}
		// 1. Best row by vertical-overlap ratio (>= 0.3), tie-broken by _ov.
		bestR := -1
		bestOv, bestOv2 := 0.3, 0.0
		for ri := range rows {
			rb := &rows[ri]
			if math.Min(b.Bottom, rb.y1)-math.Max(b.Top, rb.y0) <= 0 {
				continue // no vertical overlap
			}
			strip := pdf.TSRCell{X0: rb.stripX0, Y0: rb.y0, X1: rb.stripX1, Y1: rb.y1}
			inter := util.OverlapInter(&strip, &b)
			ov := inter / boxArea
			ov2 := 0.0
			if a := util.Area(&strip); a > 0 {
				ov2 = inter / a
			}
			// Skip unless strictly better than the current best, mirroring
			// Python's (ov, _ov) tuple ordering in find_overlapped_with_threshold.
			if !(ov > bestOv || (ov == bestOv && ov2 > bestOv2)) {
				continue
			}
			bestR, bestOv, bestOv2 = ri, ov, ov2
		}
		if bestR < 0 {
			continue
		}
		// 2. Tightest column within the matched row (no threshold).
		rb := &rows[bestR]
		bestC := -1
		bestDis := 1e9
		for _, ci := range rb.cells {
			c := &cells[ci]
			if math.Min(b.Bottom, c.Y1)-math.Max(b.Top, c.Y0) <= 0 {
				continue
			}
			if dis := tightestColumnDistance(&b, c); dis < bestDis {
				bestDis, bestC = dis, ci
			}
		}
		if bestC < 0 {
			continue
		}
		// 3. Assign, preserving the 0.85 guard for pre-filled cells only.
		target := &cells[bestC]
		if target.Text != "" && !BoxMatchesCell(*target, b, false) {
			continue
		}
		t := strings.TrimSpace(b.Text)
		if t == "" {
			continue
		}
		if cellFilled[bestC] {
			cellText[bestC] += " " + t
		} else {
			cellText[bestC] = t
			cellFilled[bestC] = true
		}
		matched++
	}

	for i := range cells {
		if cellFilled[i] {
			cells[i].Text = cellText[i]
		}
	}
	slog.Debug("fillCellTextFromBoxes done", "box_cell_matches", matched, "cells_filled", matched)
}

// tightestColumnDistance mirrors Python's find_horizontally_tightest_fit
// distance metric: the minimum of the left-edge gap, right-edge gap, and
// half the center gap. Smaller means the box sits tighter against the cell.
func tightestColumnDistance(b *pdf.TextBox, c *pdf.TSRCell) float64 {
	dis := math.Min(math.Abs(b.X0-c.X0), math.Abs(b.X1-c.X1))
	if center := math.Abs((b.X0+b.X1)-(c.X0+c.X1)) / 2; center < dis {
		dis = center
	}
	return dis
}

// BoxMatchesCell reports whether a text box's text may be assigned to a
// TSR cell. The threshold is two-stage:
//   - empty cell:        inter/boxArea >= 0.3  — matches Python's
//     find_overlapped_with_threshold default (thr=0.3), which fills cells from
//     overlapping PDF boxes uniformly.
//   - cell already has text: inter/boxArea >= 0.85 — Go-only guard, NOT in
//     Python. In the rotated-table path (table_extract.go) ocrTableCells
//     pre-fills cells with per-cell OCR text; the 0.85 bar stops a
//     weakly-overlapping detected box from corrupting/overriding that OCR
//     result. Python has no per-cell OCR at this stage, so it never raises the
//     threshold. This is a deliberate go_intentional divergence: it can drop a
//     legitimate secondary text fragment (a box overlapping 30-85%) that
//     Python would keep.
//
// FillCellTextFromBoxes uses only the 0.85 branch (cellIsEmpty=false) as the
// guard for PRE-FILLED cells; for empty cells it relies on the row/column
// selection (which already enforces >= 0.3 vertical overlap, matching Python).
// BoxMatchesCell remains the canonical "does this box match this exact cell"
// primitive and is directly unit-tested.
func BoxMatchesCell(cell pdf.TSRCell, box pdf.TextBox, cellIsEmpty bool) bool {
	inter := util.OverlapInter(&cell, &box)
	boxArea := util.Area(&box)
	if boxArea <= 0 {
		return false
	}
	if cellIsEmpty {
		return inter/boxArea >= 0.3 // Python's find_overlapped_with_threshold default
	}
	return inter/boxArea >= 0.85
}

// isCaptionBox checks if a text box is a table/figure caption,
// matching Python is_caption().  Captions should not enter table cells.
var reCaption = regexp.MustCompile(`^[图表]+[ 0-9:：]{2,}|(?i)Fig\.?\s*\d+|(?i)Figure\s+\d+|(?i)Table\s+\d+`)

func IsCaptionBox(text string, layoutType string) bool {
	if strings.Contains(layoutType, "caption") {
		return true
	}
	return reCaption.MatchString(strings.TrimSpace(text))
}

// reTableCaptionText matches text patterns that indicate a table caption
// (as opposed to a figure caption). Python is_caption uses the same set.
var reTableCaptionText = regexp.MustCompile(`^表|(?i)Table\s+\d+`)

// reFigureCaptionText matches text patterns that indicate a figure caption.
var reFigureCaptionText = regexp.MustCompile(`^图|(?i)Fig\.?\s*\d+|(?i)Figure\s+\d+`)

// captionKind returns "table" if the section is a table caption,
// "figure" if a figure caption, or "" if not a caption.
// Matches Python's is_caption check: text patterns OR layout_type containing "caption".
func CaptionKind(s pdf.Section) string {
	lt := s.LayoutType
	if lt == pdf.DLALabelTableCaption || (strings.Contains(lt, "caption") && reTableCaptionText.MatchString(strings.TrimSpace(s.Text))) {
		return pdf.LayoutTypeTable
	}
	if lt == pdf.DLALabelFigureCaption || strings.Contains(lt, "caption") {
		return pdf.LayoutTypeFigure
	}
	// DLA may label captions as "text" or other types — check text patterns.
	t := strings.TrimSpace(s.Text)
	if reTableCaptionText.MatchString(t) {
		return pdf.LayoutTypeTable
	}
	if reFigureCaptionText.MatchString(t) {
		return pdf.LayoutTypeFigure
	}
	// "图表" pattern could be either — check if isCaptionBox matches.
	if IsCaptionBox(t, "") {
		return pdf.LayoutTypeTable
	}
	return ""
}

// ── blockType: cell content classification (Python: TableStructureRecognizer.blockType) ──

// Compiled once at package init.
var blockTypePatterns = []struct {
	re   *regexp.Regexp
	kind string
}{
	// Dt (date) patterns — Python blockType lines 161-168.
	{regexp.MustCompile(`^(20|19)[0-9]{2}[年/-][0-9]{1,2}[月/-][0-9]{1,2}日*$`), "Dt"},
	{regexp.MustCompile(`^(20|19)[0-9]{2}年$`), "Dt"},
	{regexp.MustCompile(`^(20|19)[0-9]{2}[年-][0-9]{1,2}月*$`), "Dt"},
	{regexp.MustCompile(`^[0-9]{1,2}[月-][0-9]{1,2}日*$`), "Dt"},
	{regexp.MustCompile(`^第*[一二三四1-4]季度$`), "Dt"},
	{regexp.MustCompile(`^(20|19)[0-9]{2}年*[一二三四1-4]季度$`), "Dt"},
	{regexp.MustCompile(`^(20|19)[0-9]{2}[ABCDE]$`), "Dt"},
	// Nu (numeric) — Python blockType line 169.
	{regexp.MustCompile(`^[0-9.,+%/ -]+$`), "Nu"},
	// Ca (categorical) — Python blockType line 170.
	{regexp.MustCompile(`^[0-9A-Z/\._~-]+$`), "Ca"},
	// En (English) — Python blockType line 171.
	{regexp.MustCompile(`^[A-Z]*[a-z' -]+$`), "En"},
	// NE (named entity — mixed alphanumeric) — Python blockType line 172.
	{regexp.MustCompile(`^[0-9.,+-]+[0-9A-Za-z/$￥%<>（）()' -]+$`), "NE"},
	// Sg (single character) — Python blockType line 173.
	{regexp.MustCompile(`^.{1}$`), "Sg"},
}

// blockType classifies cell text into one of 9+1 types, matching Python's
// TableStructureRecognizer.blockType.  Types: Dt (date), Nu (numeric),
// Ca (categorical), En (English), NE (named entity), Sg (single char),
// Tx (short text), Lx (long text), Nr (person name), Ot (other).
func BlockType(text string) string {
	t := strings.TrimSpace(text)
	for _, p := range blockTypePatterns {
		if p.re.MatchString(t) {
			return p.kind
		}
	}
	// Token-based classification: >3 tokens, <12 → Tx, >=12 → Lx.
	// Uses simple token counting (whitespace split + individual CJK chars).
	tkn := simpleTokenCount(t)
	if tkn > 3 {
		if tkn < 12 {
			return "Tx"
		}
		return "Lx"
	}
	// Single token with POS tag "nr" → "Nr" (requires tokenizer — not available).
	// Default: "Ot" (other).
	return "Ot"
}

// simpleTokenCount estimates token count: splits on whitespace and counts
// CJK characters individually (each CJK char ≈ one token in Chinese).
func simpleTokenCount(text string) int {
	count := 0
	for _, r := range text {
		if pdf.IsCJK(r) {
			count++
		} else if r == ' ' || r == '\t' {
			// whitespace tokenizes boundaries already counted via words
		}
	}
	// Also count space-separated words.
	words := strings.FieldsSeq(text)
	for w := range words {
		if !containsCJK(w) {
			count++
		}
	}
	return count
}

func containsCJK(s string) bool {
	for _, r := range s {
		if pdf.IsCJK(r) {
			return true
		}
	}
	return false
}

// headerSetWithBlockType returns rows that should be header rows, using both
// TSR cell labels AND block-type classification.  Matches Python's
// construct_table header detection (table_structure_recognizer.py:370-384).
func HeaderSetWithBlockType(rows [][]pdf.TSRCell) map[int]bool {
	// Compute dominant block type across all cells.
	typeCounts := make(map[string]int)
	for _, row := range rows {
		for _, cell := range row {
			t := strings.TrimSpace(cell.Text)
			if t != "" {
				typeCounts[BlockType(t)]++
			}
		}
	}
	maxType := ""
	maxCount := 0
	for t, c := range typeCounts {
		if c > maxCount {
			maxType = t
			maxCount = c
		}
	}

	hdrs := make(map[int]bool)
	for ri, row := range rows {
		cnt, h := 0, 0
		for _, cell := range row {
			t := strings.TrimSpace(cell.Text)
			if t == "" {
				continue
			}
			cnt++
			bt := BlockType(t)
			// Python: if max_type == "Nu" and cell btype == "Nu" → skip
			if maxType == "Nu" && bt == "Nu" {
				continue
			}
			// Python: max_type == "Nu" and cell btype != "Nu" → header
			if maxType == "Nu" && bt != "Nu" {
				h++
			}
		}
		if cnt > 0 && float64(h)/float64(cnt) > 0.5 {
			hdrs[ri] = true
		}
	}
	// Fallback: if block-type found no headers, check for model-agnostic
	// "header" substring in cell labels (works across different TSR models).
	if len(hdrs) == 0 {
		for ri, row := range rows {
			for _, cell := range row {
				if strings.Contains(cell.Label, "header") || strings.Contains(cell.Label, "Header") {
					hdrs[ri] = true
					break
				}
			}
		}
	}
	return hdrs
}