ml_utils/extract.py

"""Spatial field extraction engine — replaces YOLO detection.

Uses OCR bounding box positions + regex to locate and extract fields.
"""

from __future__ import annotations

import re
from ml_utils.ocr import OcrResult, group_by_lines


# ── Core Spatial Helpers ────────────────────────────────────────────────
def _image_dimensions(results: list[OcrResult]) -> tuple[int, int]:
    """Estimate image dimensions from OCR bounding boxes."""
    if not results:
        return (1, 1)
    max_x = max(max(p[0] for p in r.bbox) for r in results)
    max_y = max(max(p[1] for p in r.bbox) for r in results)
    return (max(max_x, 1), max(max_y, 1))


def find_by_regex(results: list[OcrResult], pattern: str, flags: int = 0) -> list[tuple[str, OcrResult]]:
    """Find all OCR blocks whose text matches a regex pattern."""
    compiled = re.compile(pattern, flags)
    matches: list[tuple[str, OcrResult]] = []
    for r in results:
        m = compiled.search(r.text)
        if m:
            matches.append((m.group(0), r))
    return matches


def find_keyword(results: list[OcrResult], keywords: list[str]) -> OcrResult | None:
    """Find first OCR block containing any of the keywords (case-insensitive)."""
    lower_kw = [k.lower() for k in keywords]
    for r in results:
        text_lower = r.text.lower()
        for kw in lower_kw:
            if kw in text_lower:
                return r
    return None


def find_value_near_label(
    results: list[OcrResult],
    label_keywords: list[str],
    direction: str = "right_or_below",
    max_distance_pct: float = 0.15,
) -> str | None:
    """Find the text block nearest to a label keyword.

    direction: 'right', 'below', 'right_or_below'
    max_distance_pct: max distance as fraction of image dimension
    """
    label_block = find_keyword(results, label_keywords)
    if label_block is None:
        return None

    img_w, img_h = _image_dimensions(results)
    lx1, ly1, lx2, ly2 = label_block.rect
    max_dx = img_w * max_distance_pct
    max_dy = img_h * max_distance_pct

    candidates: list[tuple[float, OcrResult]] = []

    for r in results:
        if r is label_block:
            continue
        rx1, ry1, rx2, ry2 = r.rect

        if direction in ("right", "right_or_below"):
            # Block is to the right and roughly same vertical position
            if rx1 >= lx2 - 10 and abs(r.center_y - label_block.center_y) < max_dy:
                dist = rx1 - lx2
                if dist < max_dx * 3:  # more lenient horizontally
                    candidates.append((dist, r))

        if direction in ("below", "right_or_below"):
            # Block is below and roughly same horizontal position
            if ry1 >= ly2 - 10 and abs(r.center_x - label_block.center_x) < max_dx * 2:
                dist = ry1 - ly2
                if dist < max_dy * 2:
                    candidates.append((dist + 10000, r))  # prefer right over below

    if not candidates:
        # Fallback: check if label text contains value after colon
        colon_match = re.search(r'[:\-]\s*(.+)', label_block.text)
        if colon_match:
            value = colon_match.group(1).strip()
            if value:
                return value
        return None

    candidates.sort(key=lambda x: x[0])
    return candidates[0][1].text


def get_text_in_region(
    results: list[OcrResult],
    y_start_pct: float = 0.0,
    y_end_pct: float = 1.0,
    x_start_pct: float = 0.0,
    x_end_pct: float = 1.0,
) -> str:
    """Get all text in a rectangular region (percentages of image dimensions)."""
    img_w, img_h = _image_dimensions(results)
    y_start = img_h * y_start_pct
    y_end = img_h * y_end_pct
    x_start = img_w * x_start_pct
    x_end = img_w * x_end_pct

    region_results = [
        r for r in results
        if y_start <= r.center_y <= y_end and x_start <= r.center_x <= x_end
    ]
    lines = group_by_lines(region_results)
    return "\n".join(" ".join(r.text for r in line) for line in lines)


def get_all_text_blocks_in_region(
    results: list[OcrResult],
    y_start_pct: float = 0.0,
    y_end_pct: float = 1.0,
    x_start_pct: float = 0.0,
    x_end_pct: float = 1.0,
) -> list[OcrResult]:
    """Get all OCR blocks in a rectangular region."""
    img_w, img_h = _image_dimensions(results)
    y_start = img_h * y_start_pct
    y_end = img_h * y_end_pct
    x_start = img_w * x_start_pct
    x_end = img_w * x_end_pct

    return [
        r for r in results
        if y_start <= r.center_y <= y_end and x_start <= r.center_x <= x_end
    ]


def get_nearby_text(
    results: list[OcrResult],
    anchor: OcrResult,
    direction: str = "below",
    max_blocks: int = 5,
    max_distance_pct: float = 0.2,
) -> list[OcrResult]:
    """Get text blocks near an anchor block in a given direction."""
    img_w, img_h = _image_dimensions(results)
    ax1, ay1, ax2, ay2 = anchor.rect
    max_dist = img_h * max_distance_pct if direction in ("below", "above") else img_w * max_distance_pct

    candidates: list[tuple[float, OcrResult]] = []
    for r in results:
        if r is anchor:
            continue
        rx1, ry1, rx2, ry2 = r.rect

        if direction == "below" and ry1 >= ay2 - 5:
            dist = ry1 - ay2
            if dist < max_dist:
                candidates.append((dist, r))
        elif direction == "right" and rx1 >= ax2 - 5:
            dist = rx1 - ax2
            if dist < max_dist and abs(r.center_y - anchor.center_y) < anchor.height * 1.5:
                candidates.append((dist, r))

    candidates.sort(key=lambda x: x[0])
    return [c[1] for c in candidates[:max_blocks]]


# ── Document-Specific Extraction Router ─────────────────────────────────
def extract_fields(ocr_results: list[OcrResult], doc_type: str) -> dict:
    """Route to the correct document-specific extractor."""
    if doc_type == "aadhaar":
        from ml_utils.validators.aadhaar import parse_aadhaar_fields
        return parse_aadhaar_fields(ocr_results)
    elif doc_type == "pan":
        from ml_utils.validators.pan import parse_pan_fields
        return parse_pan_fields(ocr_results)
    elif doc_type == "caste":
        from ml_utils.validators.caste import parse_caste_fields_from_ocr
        return parse_caste_fields_from_ocr(ocr_results)
    elif doc_type == "experience":
        from ml_utils.validators.experience import parse_experience_fields_from_ocr
        return parse_experience_fields_from_ocr(ocr_results)
    elif doc_type == "education":
        from ml_utils.validators.education import parse_education_fields
        return parse_education_fields(ocr_results)
    elif doc_type == "resume":
        from ml_utils.validators.resume import parse_resume_fields
        return parse_resume_fields(ocr_results)
    else:
        from ml_utils.validators.general import parse_general_fields
        return parse_general_fields(ocr_results)