src/docflow/extractor.py

"""Field extraction from OCR text using regex, heuristics, line parsing, and NER."""

from __future__ import annotations

import re
from statistics import mean

from docflow.models import ExtractedInvoice, LineItem
from docflow.ner import extract_organizations, find_tax_ids

_PERSIAN_DIGITS = str.maketrans("۰۱۲۳۴۵۶۷۸۹", "0123456789")
_ARABIC_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789")

FIELD_ALIASES: dict[str, list[str]] = {
    "vendor_name": ["فروشنده", "صادرکننده", "vendor", "seller", "company"],
    "vendor_tax_id": ["شناسه ملی", "کد اقتصادی", "شماره اقتصادی", "tax id", "vat no", "economic code"],
    "invoice_number": ["شماره فاکتور", "ش فاکتور", "invoice no", "invoice #", "invoice number"],
    "invoice_date_jalali": ["تاریخ", "date"],
    "buyer_name": ["خریدار", "مشتری", "buyer", "customer"],
    "subtotal": ["جمع بدون مالیات", "subtotal", "net amount"],
    "tax_amount": ["مالیات بر ارزش افزوده", "ارزش افزوده", "vat", "tax"],
    "total_amount": ["جمع کل", "مبلغ قابل پرداخت", "grand total", "total amount"],
}

LINE_ITEM_PATTERN = re.compile(
    r"^(.+?)\s{2,}(\d+(?:\.\d+)?)\s{2,}([\d,\.]+)\s{2,}([\d,\.]+)\s*$"
)
LINE_ITEM_PATTERN_ALT = re.compile(
    r"^(.+?)\s+(\d+(?:\.\d+)?)\s+([\d,\.]+)\s+([\d,\.]+)\s*$"
)


def normalize_text(text: str) -> str:
    text = text.translate(_PERSIAN_DIGITS).translate(_ARABIC_DIGITS)
    return text.replace("،", ",").replace("٬", ",").replace(" : ", ": ")


def _parse_amount(value: str | None) -> float | None:
    if not value:
        return None
    cleaned = re.sub(r"[^\d.]", "", value.replace(",", ""))
    try:
        return float(cleaned)
    except ValueError:
        return None


def _line_key_value(lines: list[str]) -> dict[str, str]:
    """Parse 'Label: Value' patterns line by line."""
    result: dict[str, str] = {}
    for line in lines:
        line = line.strip()
        if not line:
            continue
        for field, aliases in FIELD_ALIASES.items():
            for alias in aliases:
                pattern = rf"(?i)^{re.escape(alias)}\s*[:\-]\s*(.+)$"
                match = re.match(pattern, line)
                if match:
                    result[field] = match.group(1).strip()
                    break
    return result


def _regex_fallback(text: str, field: str) -> str | None:
    patterns = {
        "vendor_name": [r"(?:فروشنده|صادرکننده|شرکت)\s*[:\-]\s*(.+)", r"(?:vendor|seller)\s*[:\-]\s*(.+)"],
        "vendor_tax_id": [r"(?:شناسه\s*ملی|کد\s*اقتصادی)\s*[:\-]\s*(\d{10,14})", r"(?:tax\s*id)\s*[:\-]\s*(\d{10,14})"],
        "invoice_number": [r"(?:شماره\s*فاکتور|invoice\s*(?:no|#)?)\s*[:\-]\s*([A-Za-z0-9\-/]+)"],
        "invoice_date_jalali": [r"(?:تاریخ|date)\s*[:\-]\s*(\d{4}[/\-]\d{1,2}[/\-]\d{1,2})", r"(\d{4}/\d{2}/\d{2})"],
        "buyer_name": [r"(?:خریدار|مشتری|buyer)\s*[:\-]\s*(.+)"],
        "subtotal": [r"(?:جمع\s*بدون\s*مالیات|subtotal)\s*[:\-]\s*([\d,\.]+)"],
        "tax_amount": [r"(?:مالیات|vat|tax)\s*[:\-]\s*([\d,\.]+)"],
        "total_amount": [r"(?:جمع\s*کل|grand\s*total|total\s*amount)\s*[:\-]\s*([\d,\.]+)"],
    }
    for pattern in patterns.get(field, []):
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            return match.group(1).strip()
    return None


def _extract_line_items(text: str) -> list[LineItem]:
    items: list[LineItem] = []
    for line in text.splitlines():
        line = line.strip()
        for pattern in (LINE_ITEM_PATTERN, LINE_ITEM_PATTERN_ALT):
            match = pattern.match(line)
            if match:
                desc, qty, unit, total = match.groups()
                if desc.lower() in ("description", "شرح کالا", "item"):
                    continue
                items.append(
                    LineItem(
                        description=desc.strip(),
                        quantity=float(qty),
                        unit_price=_parse_amount(unit),
                        total=_parse_amount(total),
                        confidence=0.82,
                    )
                )
                break
    return items


def _estimate_confidence(invoice: ExtractedInvoice, ocr_blocks: list[dict]) -> float:
    key_fields = [invoice.vendor_name, invoice.invoice_number, invoice.total_amount, invoice.invoice_date_jalali]
    filled = sum(1 for f in key_fields if f) / len(key_fields)
    ocr_score = mean(b["confidence"] for b in ocr_blocks) if ocr_blocks else 0.7
    return round(0.6 * filled + 0.4 * ocr_score, 3)


def extract_fields(raw_text: str, ocr_blocks: list[dict] | None = None) -> ExtractedInvoice:
    text = normalize_text(raw_text)
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    kv = _line_key_value(lines)
    ocr_blocks = ocr_blocks or []

    vendor = kv.get("vendor_name") or _regex_fallback(text, "vendor_name")
    if not vendor:
        orgs = extract_organizations(text)
        vendor = orgs[0] if orgs else None

    tax_id = kv.get("vendor_tax_id") or _regex_fallback(text, "vendor_tax_id")
    if not tax_id:
        ids = find_tax_ids(text)
        tax_id = ids[0] if ids else None

    invoice = ExtractedInvoice(
        vendor_name=vendor,
        vendor_tax_id=tax_id,
        invoice_number=kv.get("invoice_number") or _regex_fallback(text, "invoice_number"),
        invoice_date_jalali=kv.get("invoice_date_jalali") or _regex_fallback(text, "invoice_date_jalali"),
        buyer_name=kv.get("buyer_name") or _regex_fallback(text, "buyer_name"),
        subtotal=_parse_amount(kv.get("subtotal") or _regex_fallback(text, "subtotal")),
        tax_amount=_parse_amount(kv.get("tax_amount") or _regex_fallback(text, "tax_amount")),
        total_amount=_parse_amount(kv.get("total_amount") or _regex_fallback(text, "total_amount")),
        line_items=_extract_line_items(text),
        raw_text=raw_text,
        extraction_method="hybrid_ocr_ner_regex",
    )

    lower = text.lower()
    if "تومان" in lower or "toman" in lower:
        invoice.currency = "IRT"
    elif "usd" in lower or "دلار" in text:
        invoice.currency = "USD"

    invoice.confidence = _estimate_confidence(invoice, ocr_blocks)
    return invoice