core/region_classifier.py

"""
Semantic & Geometric Region Classifier.
Accurately categorizes detected bounding boxes into the 10 standardized categories:
1. Title/Header
2. Paragraph
3. Text
4. Table
5. Table cell
6. Key-Value
7. Number/Price
8. Image
9. Footer
10. Other
"""

import re
from typing import List, Optional, Tuple, Dict, Any
from core.models import RegionType

# Regex patterns for fast, robust semantic detection
PRICE_CURRENCY_PATTERN = re.compile(
    r"^[\$€£¥₫]?\s*[\d]{1,3}(?:[.,]\d{3})*(?:[.,]\d{1,2})?\s*[\$€£¥₫%]?$|"
    r"^(?:total|subtotal|tax|amount|due|balance|change|cash|visa|mastercard|tip|vat|discount)\s*[:#\-=]?\s*[\$€£¥₫]?\s*[\d]+[.,]?\d*$",
    re.IGNORECASE
)

NUMBER_ONLY_PATTERN = re.compile(r"^#?\s*\d+(?:[.,]\d+)*\s*%?$")

KEY_VALUE_PATTERN = re.compile(
    r"^(?:[A-Za-z0-9\s/_\-]{2,30})\s*[:=]\s*(?:.+)$|"
    r"^(?:date|time|cashier|order|invoice|tel|phone|table|server|receipt|card|auth|ref|account|cust|item|addr|address)\s*[:#\-]\s*(?:.+)$",
    re.IGNORECASE
)

FOOTER_KEYWORDS = [
    "thank you", "thanks", "visit again", "have a nice day", "powered by",
    "customer copy", "merchant copy", "terms and conditions", "page ", "www.", ".com",
    "all rights reserved", "copyright", "cảm ơn", "hen gap lai", "hẹn gặp lại"
]

HEADER_KEYWORDS = [
    "receipt", "invoice", "tax invoice", "sales receipt", "bill", "statement",
    "hóa đơn", "phiếu thanh toán", "phiếu thu", "menu", "store", "supermarket",
    "restaurant", "official receipt"
]


def classify_region(
    text: str,
    box: List[int],
    image_width: int,
    image_height: int,
    native_tag: Optional[str] = None
) -> str:
    """
    Classifies a bounding box region into one of the 10 standardized types.
    Considers native model tag (if present), spatial geometry, and textual semantics.
    """
    cleaned_text = (text or "").strip()
    x1, y1, x2, y2 = box
    box_w = max(1, x2 - x1)
    box_h = max(1, y2 - y1)
    center_y = (y1 + y2) / 2.0
    relative_y = center_y / max(1.0, float(image_height))

    # 1. Native Model Tag overrides if explicit
    if native_tag:
        tag_lower = native_tag.lower()
        if "cell" in tag_lower:
            return RegionType.TABLE_CELL.value
        if "table" in tag_lower or "grid" in tag_lower:
            return RegionType.TABLE.value
        if "figure" in tag_lower or "image" in tag_lower or "logo" in tag_lower or "seal" in tag_lower or "stamp" in tag_lower:
            return RegionType.IMAGE.value
        if "header" in tag_lower or "title" in tag_lower:
            return RegionType.TITLE_HEADER.value
        if "footer" in tag_lower:
            return RegionType.FOOTER.value

    if not cleaned_text:
        # Empty text box is likely an image/logo/separator or other
        if box_w > 50 and box_h > 50:
            return RegionType.IMAGE.value
        return RegionType.OTHER.value

    # 2. Markdown / HTML tag hints
    if "<table" in cleaned_text.lower() or ("|" in cleaned_text and "-|-" in cleaned_text):
        return RegionType.TABLE.value
    if cleaned_text.startswith("# ") or cleaned_text.startswith("## "):
        return RegionType.TITLE_HEADER.value

    # 3. Number / Price check (highest priority for amounts)
    if PRICE_CURRENCY_PATTERN.match(cleaned_text) or (NUMBER_ONLY_PATTERN.match(cleaned_text) and len(cleaned_text) < 15):
        return RegionType.NUMBER_PRICE.value

    # 4. Key-Value check
    if KEY_VALUE_PATTERN.match(cleaned_text) and ":" in cleaned_text:
        return RegionType.KEY_VALUE.value

    # 5. Footer check
    if relative_y > 0.88 or any(kw in cleaned_text.lower() for kw in FOOTER_KEYWORDS):
        if relative_y > 0.65:
            return RegionType.FOOTER.value

    # 6. Title / Header check
    if relative_y < 0.25 and ":" not in cleaned_text:
        is_header_kw = any(kw in cleaned_text.lower() for kw in HEADER_KEYWORDS)
        is_prominent_title = (cleaned_text.isupper() and len(cleaned_text) < 50) or is_header_kw or (box_h > 35)
        if is_prominent_title:
            return RegionType.TITLE_HEADER.value

    # 7. Paragraph vs Single Text Line
    word_count = len(cleaned_text.split())
    if word_count >= 12 or box_h > 50 or "\n" in cleaned_text:
        return RegionType.PARAGRAPH.value

    # Default to general Text
    return RegionType.TEXT.value


def compute_region_summary(regions: List[Any]) -> Dict[str, int]:
    """Computes count of regions per category."""
    counts = {rt.value: 0 for rt in RegionType}
    for r in regions:
        rtype = getattr(r, "region_type", None) or (r.get("region_type") if isinstance(r, dict) else RegionType.TEXT.value)
        if rtype in counts:
            counts[rtype] += 1
        else:
            counts[RegionType.OTHER.value] += 1
    return counts