extractors.py

"""Keyword and structured-field extraction for OCR'd document text.

Domain keyword sets flag "notification" events (things a reviewer would want
alerted on), and field patterns pull structured values (IDs, dates, amounts)
out of the surrounding text so results can be dropped straight into a CSV.
"""
import re

KEYWORD_SETS = {
    "construction": [
        "change order", "delay notice", "punch list", "non-conformance",
        "rfi", "submittal rejected", "back charge", "backcharge",
        "stop work", "safety violation", "material shortage",
        "schedule slip", "cost overrun", "field order", "notice to proceed",
    ],
    "compliance": [
        "non-compliant", "noncompliance", "violation", "deficiency",
        "corrective action", "audit finding", "expired", "expiration",
        "recall", "citation", "penalty", "out of compliance",
        "failed inspection", "not approved",
    ],
    "finance": [
        "past due", "overdue", "payment rejected", "invoice disputed",
        "credit hold", "chargeback", "insufficient funds", "write-off",
        "collections", "declined", "late fee", "short paid",
    ],
    "general": [
        "urgent", "immediate action required", "cancelled", "rejected",
        "approved", "pending approval", "revised", "on hold",
    ],
}

FIELD_PATTERNS = {
    "dates": re.compile(
        r"\b(?:\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4}-\d{2}-\d{2}|"
        r"(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\.?\s+"
        r"\d{1,2},?\s+\d{4})\b",
        re.IGNORECASE,
    ),
    "amounts": re.compile(
        r"\$\s?\d[\d,]*\.?\d{0,2}\b|\bUSD\s?\d[\d,]*\.?\d{0,2}\b",
        re.IGNORECASE,
    ),
    "percentages": re.compile(r"\b\d{1,3}(?:\.\d+)?\s?%"),
    "po_numbers": re.compile(
        r"\bP\.?\s?O\.?[ \t]*#?[ \t]*[:\-]?[ \t]*((?=[A-Z0-9-]*\d)[A-Z0-9\-]{3,15})\b",
        re.IGNORECASE,
    ),
    "sap_doc_numbers": re.compile(
        r"\bSAP[ \t]*(?:Doc(?:ument)?\.?|#)?[ \t]*[:\-]?[ \t]*(\d{8,10})\b|\b(4[5-8]\d{8})\b",
        re.IGNORECASE,
    ),
    "invoice_numbers": re.compile(
        r"\bINV(?:OICE)?\.?[ \t]*#?[ \t]*[:\-]?[ \t]*((?=[A-Z0-9-]*\d)[A-Z0-9\-]{3,15})\b",
        re.IGNORECASE,
    ),
    "cert_codes": re.compile(
        r"\b(ISO\s?\d{3,5}(?:-\d+)?|OSHA\s?\d{3,4}(?:\.\d+)?|"
        r"ASTM\s?[A-Z]?\d{2,4}(?:-\d{2})?|ANSI\s?[A-Z]?\d+(?:\.\d+)?)\b",
        re.IGNORECASE,
    ),
}

# Which standard field patterns are relevant to each domain. Extraction is
# scoped to the categories a user actually selects rather than running every
# pattern against every document ("blanket" extraction) - a compliance
# certificate has no business generating empty invoice_number columns.
CATEGORY_FIELD_MAP = {
    "construction": ["po_numbers", "sap_doc_numbers", "dates", "amounts"],
    "compliance": ["cert_codes", "dates"],
    "finance": ["invoice_numbers", "amounts", "dates", "percentages"],
    "general": ["dates"],
}

CONTEXT_WINDOW = 60
MAX_CUSTOM_VALUE_LEN = 80


def _dedupe(values):
    seen = []
    for v in values:
        v = v.strip()
        if v and v not in seen:
            seen.append(v)
    return seen


def fields_for_categories(categories):
    """Union of standard field names relevant to the given categories."""
    names = []
    for cat in categories:
        for name in CATEGORY_FIELD_MAP.get(cat, []):
            if name not in names:
                names.append(name)
    return names


def extract_fields(text, categories):
    """Pull the standard fields relevant to `categories` out of a page of text.

    Returns {field_name: [values]} scoped to only the patterns those
    categories map to - see CATEGORY_FIELD_MAP.
    """
    fields = {}
    for name in fields_for_categories(categories):
        pattern = FIELD_PATTERNS[name]
        matches = pattern.findall(text)
        flat = []
        for m in matches:
            if isinstance(m, tuple):
                flat.extend(g for g in m if g)
            else:
                flat.append(m)
        fields[name] = _dedupe(flat)
    return fields


def _postprocess_custom_value(label, raw_value):
    raw_value = raw_value.strip()[:MAX_CUSTOM_VALUE_LEN]
    if "date" in label:
        m = FIELD_PATTERNS["dates"].search(raw_value)
        if m:
            return m.group(0)
    if any(k in label for k in ("weight", "amount", "qty", "quantity", "total", "count", "price")):
        m = re.match(r"[\d,]+\.?\d*\s?[a-zA-Z%]*", raw_value)
        if m and m.group(0).strip():
            return m.group(0).strip()
    # Generic fallback: a handful of tokens, cut before an obvious next column
    # (two-or-more spaces, which OCR/tables often use as a column separator).
    trimmed = re.split(r"\s{2,}", raw_value)[0]
    m = re.match(r"\S+(?:\s\S+){0,4}", trimmed)
    return m.group(0).strip() if m else trimmed.strip()


def extract_custom_fields(text, field_labels):
    """Generic label->value extraction for user-defined fields.

    For labels like "reference number" or "gross weight", finds the label in
    the text and captures whatever follows it on the same line (or the next
    line, if the label sits alone on its own line - common in OCR'd forms).
    Returns {label: [values]}.
    """
    results = {label: [] for label in field_labels if label.strip()}
    if not results:
        return results

    lines = text.splitlines()
    for i, line in enumerate(lines):
        lower_line = line.lower()
        for label in results:
            label_clean = label.strip().lower()
            idx = lower_line.find(label_clean)
            if idx == -1:
                continue
            after = line[idx + len(label_clean):]
            after = re.sub(r"^[\s:.\-–—]+", "", after).strip()
            if not after and i + 1 < len(lines):
                after = lines[i + 1].strip()
            if after:
                value = _postprocess_custom_value(label_clean, after)
                if value and value not in results[label]:
                    results[label].append(value)
    return results


def find_notifications(text, categories, custom_keywords=None):
    """Scan text for keyword hits across the given categories.

    Returns a list of dicts: {category, keyword, context}.
    """
    hits = []
    keywords_by_category = {}
    for cat in categories:
        keywords_by_category[cat] = list(KEYWORD_SETS.get(cat, []))
    if custom_keywords:
        keywords_by_category["custom"] = list(custom_keywords)

    lower_text = text.lower()
    for category, keywords in keywords_by_category.items():
        for kw in keywords:
            kw_clean = kw.strip().lower()
            if not kw_clean:
                continue
            start = 0
            while True:
                idx = lower_text.find(kw_clean, start)
                if idx == -1:
                    break
                ctx_start = max(0, idx - CONTEXT_WINDOW)
                ctx_end = min(len(text), idx + len(kw_clean) + CONTEXT_WINDOW)
                context = text[ctx_start:ctx_end].replace("\n", " ").strip()
                hits.append({
                    "category": category,
                    "keyword": kw.strip(),
                    "context": context,
                })
                start = idx + len(kw_clean)
    return hits