PaperPod/paperpod/ocr/extract.py

"""OCR fallback naming: pull vendor / date / total out of a document crop.

Used when no spoken description is available (module 4, not built yet).
Tesseract (via pytesseract) does the character recognition; everything else
here is regex/heuristics tuned for North American receipts and letters.
"""

from __future__ import annotations

import re
import statistics
from dataclasses import dataclass
from datetime import datetime

import cv2
import numpy as np
import pytesseract

MONTHS = {
    "jan": 1, "january": 1, "feb": 2, "february": 2, "mar": 3, "march": 3,
    "apr": 4, "april": 4, "may": 5, "jun": 6, "june": 6, "jul": 7, "july": 7,
    "aug": 8, "august": 8, "sep": 9, "sept": 9, "september": 9, "oct": 10,
    "october": 10, "nov": 11, "november": 11, "dec": 12, "december": 12,
}

# Ordered most- to least-specific; first match wins. Slash dates tolerate a
# single OCR-injected space after each separator ("01/ 31/ 2025" on TD
# dot-matrix statements).
_DATE_PATTERNS: list[tuple[re.Pattern, str]] = [
    # Statement period end ("PERIOD: FROM 01/01/2025 TO: 01/31/2025") — the
    # canonical document date for a bank statement, and checked first so the
    # first transaction row's date doesn't win.
    (
        re.compile(
            r"\b(?:TO|AU)\s*:?\s+(\d{1,2})\s?/\s?(\d{1,2})\s?/\s?(\d{4})\b",
            re.IGNORECASE,
        ),
        "mdy",
    ),
    (re.compile(r"\b(\d{4})[-/](\d{1,2})[-/](\d{1,2})\b"), "ymd"),
    (
        re.compile(
            r"\b(jan|feb|mar|apr|may|jun|jul|aug|sep|sept|oct|nov|dec)"
            r"[a-z]*\.?\s+(\d{1,2})[.,]?\s+(\d{4})\b",
            re.IGNORECASE,
        ),
        "month_name_dmy",
    ),
    # Assumes North American MM/DD/YYYY (matches Home Depot / Metro / Petro-Canada
    # style receipts this app was designed around). Revisit if you record
    # receipts using DD/MM/YYYY formatting.
    (re.compile(r"\b(\d{1,2})\s?/\s?(\d{1,2})\s?/\s?(\d{4})\b"), "mdy"),
    (re.compile(r"\b(\d{1,2})-(\d{1,2})-(\d{4})\b"), "mdy"),
]

# OCR misreads the digit 0 as the letter O in dotted print ("O1/ 31/2025").
_O_BEFORE_DIGIT_RE = re.compile(r"[Oo](?=\d)")
_O_AFTER_DIGIT_RE = re.compile(r"(?<=\d)[Oo]")

_TOTAL_LINE_RE = re.compile(r"\btotal\b", re.IGNORECASE)
_MONEY_RE = re.compile(r"\$?\s*(\d{1,4}\.\d{2})\b")

# Line items: "DESC 12.34" / "1234567 DESC 12.34 H" (Costco tax code).
# Totals / tender / card boilerplate intentionally excluded so callers can
# itemize SKUs for split booking (kids vs home, etc.) in the taxes pipeline.
_ITEM_LINE_RE = re.compile(
    r"^(?:(?:\d{4,14})\s+)?(.+?)\s+\$?\s*(\d{1,4}\.\d{2})\s*[A-Z]?\s*$",
    re.IGNORECASE,
)
_ITEM_SKIP_RE = re.compile(
    r"\b("
    r"sub\s*total|subtotal|total|hst|gst|pst|qst|tax|change|cash|debit|"
    r"mastercard|visa|amex|interac|approved|thank|balance|save|"
    r"account|card\s*number|auth|approval|reference|invoice|"
    r"customer\s*copy|merchant\s*copy|amount\s*:|date/?time|"
    r"no\s*return|signature|member|terminal|cashier"
    r")\b",
    re.IGNORECASE,
)

# HH:MM(:SS)? 24-hour clock, as printed near the transaction date on most
# receipts (e.g. "03/20/23  15:25:35"). Distinct enough from other
# colon-free receipt numbers (barcodes, reference IDs) that the first match
# in the document is reliably the transaction time.
_TIME_RE = re.compile(r"\b([01]?\d|2[0-3]):([0-5]\d)(?::[0-5]\d)?")

# Common Canadian tax-slip codes, longest/most-specific first so "T4A" and
# "T5008" match before the more generic "T4"/"T5" would. Receipts and
# letters won't match any of these, so this only ever affects naming for
# actual government forms.
_FORM_CODES = ["T5008", "T4A", "T4E", "T2202", "RL-1", "RL-2", "T5", "T4", "T3", "T1"]
_FORM_CODE_RE = re.compile(
    r"\b(" + "|".join(re.escape(c) for c in _FORM_CODES) + r")\b"
)

# Boilerplate that shows up near the top of CRA slips and would otherwise be
# mistaken for the employer/payer name by a naive "first all-caps line" scan.
_FORM_BOILERPLATE_RE = re.compile(
    r"CANADA REVENUE|AGENCE DU REVENU|AGENCY|PROTECTED B|PROTEG|STATEMENT OF"
    r"|ETAT DU|ETAT DE|SUMMARY OF|EMPLOYER|EMPLOYEE|EMPLOYE|PAYER|PAYOUR|RECIPIENT"
    r"|YEAR|ANNEE|BOX \d|ACCOUNT NO|IDENTIFICATION|REPORT CODE|ORIGINAL|AMENDED"
    r"|DU CANADA",
    re.IGNORECASE,
)
# 2-5 consecutive ALL-CAPS words, e.g. "YORK UNIVERSITY", "AMC ENTERTAINMENT
# HOLDINGS INC". Searched as a substring (not a full-line match) since OCR
# on boxed forms often merges an organization name with adjacent column
# text ("YORK UNIVERSITY Year 2023 Statement of...") on one line.
_ORG_NAME_RE = re.compile(r"\b[A-Z][A-Z&.']{1,20}(?:\s+[A-Z][A-Z0-9&.',-]{1,20}){1,4}\b")

# Known retail chains — searched anywhere in OCR text when the vendor-line
# heuristic fails (common on narrow thermal receipts with hands in frame).
_STORE_PATTERNS: list[tuple[re.Pattern, str]] = [
    (re.compile(r"COSTCO|WHOLE?\s*SALE", re.IGNORECASE), "Costco"),
    (re.compile(r"WALMART", re.IGNORECASE), "Walmart"),
    (re.compile(r"HOME\s+DEPOT", re.IGNORECASE), "Home Depot"),
    (re.compile(r"\bMETRO\b", re.IGNORECASE), "Metro"),
    (re.compile(r"PETRO[\s-]*CANADA", re.IGNORECASE), "Petro-Canada"),
]

_TAX_YEAR_RE = re.compile(
    r"\b(?:year|ann[eé]e)\s*[:\s]?\s*(\d{4})\b", re.IGNORECASE
)
_TAX_YEAR_TITLE_RE = re.compile(
    r"\b(?:SUMMARY|STATEMENT|DISPOSITIONS|REMUNERATION)\b[^\n]{0,40}\b(20\d{2})\b",
    re.IGNORECASE,
)

# "PAGE 3 OF 17" / "PAGE 3 DE 17" (bilingual Canadian statements). Dot-matrix
# print OCRs sloppily, so the digits also admit the classic confusables
# ("PAGE ll OF 17" for 11, "0F" for OF) and stray underscores/punctuation.
_PAGE_MARKER_RE = re.compile(
    r"\bPAGE\s*[:.,]?\s*([0-9lIiO]{1,3})\s*(?:[O0Q]F|DE)[_.,:]?\s+([0-9lIiO]{1,3})\b",
    re.IGNORECASE,
)
_DIGIT_CONFUSABLES = str.maketrans({"l": "1", "I": "1", "i": "1", "O": "0", "o": "0"})

# Printed document titles worth carrying into the filename ("what is this?"),
# checked in order. Bank statements all say "TD Canada Trust" in the corner;
# the title line is what actually distinguishes a transaction history from a
# direct-deposit form.
_DOC_TITLE_PATTERNS: list[tuple[re.Pattern, str]] = [
    (re.compile(r"DIRECT\s+DEPOSIT", re.IGNORECASE), "direct_deposit"),
    (
        re.compile(r"TRANSACTION\s+HISTORY|HISTORIQUE\s+DES\s+OP", re.IGNORECASE),
        "transaction_history",
    ),
    (re.compile(r"DEPOSIT\s+ACCOUNT\s+HISTORY", re.IGNORECASE), "transaction_history"),
    (re.compile(r"ACCOUNT\s+STATEMENT|RELEV[EÉ]\s+DE\s+COMPTE", re.IGNORECASE), "statement"),
    (re.compile(r"\bINVOICE\b|\bFACTURE\b", re.IGNORECASE), "invoice"),
    (re.compile(r"NOTICE\s+OF\s+ASSESSMENT", re.IGNORECASE), "notice_of_assessment"),
    (re.compile(r"VOID\s+CH(?:EQUE|ECK)", re.IGNORECASE), "void_cheque"),
]

# Letters only (Latin + accented) — used to judge whether a vendor token is
# a real word after ignoring OCR-injected punctuation ("PETRO-—CANADA").
_LETTERS_RE = re.compile(r"[^A-Za-zÀ-ÖØ-öø-ÿ]+")
_VOWELS = "aeiouyàâäéèêëîïôöùûü"


def extract_page_marker(text: str) -> tuple[int, int] | None:
    """(page_number, page_total) from a printed 'PAGE X OF Y', if present."""
    match = _PAGE_MARKER_RE.search(text)
    if not match:
        return None
    try:
        page = int(match.group(1).translate(_DIGIT_CONFUSABLES))
        total = int(match.group(2).translate(_DIGIT_CONFUSABLES))
    except ValueError:
        return None
    if not (1 <= page <= total <= 500):
        return None
    return page, total


def extract_doc_title(text: str) -> str | None:
    """Recognized printed document title (as a filename-ready slug)."""
    for pattern, slug in _DOC_TITLE_PATTERNS:
        if pattern.search(text):
            return slug
    return None


# TD / credit-union letterheads print the account holder as SHORTNAME.
# Also match "Account name:" / "Customer:" / name-on-file style labels.
_PARTY_RE = re.compile(
    r"\b(?:SHORT\s*NAME|ACCOUNT\s*NAME|CUSTOMER(?:\s*NAME)?|NAME)\s*:?\s*"
    r"([A-Z][A-Z0-9 .'-]{1,40})",
    re.IGNORECASE,
)


def extract_party(text: str) -> str | None:
    """Account holder / shortname printed on the document, if labeled."""
    match = _PARTY_RE.search(text)
    if not match:
        return None
    name = re.sub(r"\s+", " ", match.group(1)).strip(" .,-")
    # Reject labels that ate the next field ("LEVIT I PERIOD").
    name = re.split(r"\b(?:PERIOD|BR|ACCOUNT|FROM|TO|PAGE)\b", name, maxsplit=1)[0].strip()
    if len(name) < 2:
        return None
    return name


def clean_vendor(vendor: str | None) -> str | None:
    """Strip OCR junk from a vendor line; reject it entirely if it's gibberish.

    Tesseract routinely bolts noise onto real letterhead text: "| ») Transaction
    History i", "7D TD Canada Trust", "+7 TD Canada Trus". Leading/trailing
    tokens that aren't real words get trimmed; if fewer than half the remaining
    tokens look wordlike ("BH tc rn te So Gee EA Nie..."), the whole line is
    treated as noise and dropped so naming can fall back to something better.
    """
    if not vendor:
        return None
    tokens = vendor.split()

    def is_word(tok: str) -> bool:
        # Judge on letters only, so hyphens/OCR dashes inside real names
        # ("PETRO-—CANADA") don't disqualify the token.
        core = _LETTERS_RE.sub("", tok)
        return len(core) >= 3 and any(c in _VOWELS for c in core.lower())

    def is_acronym(tok: str) -> bool:
        # Pure 2-5 uppercase letters ("TD", "CRA") — "7D" or "»)" don't count.
        return tok.isalpha() and tok.isupper() and 2 <= len(tok) <= 5

    def keepable(tok: str) -> bool:
        return is_word(tok) or is_acronym(tok)

    while tokens and not keepable(tokens[0]):
        tokens.pop(0)
    while tokens and not keepable(tokens[-1]):
        tokens.pop()
    if not tokens:
        return None
    words = sum(1 for t in tokens if is_word(t))
    if words < max(1, len(tokens) / 2):
        return None
    return " ".join(tokens)


# Real CRA slips always carry recognizable boilerplate near the code. Short
# codes like "T4" otherwise false-positive on garbled OCR of unrelated
# documents (a misread "LEVIT I" once produced a confident "T4").
_FORM_CONTEXT_RE = re.compile(
    r"STATEMENT OF|ETAT D|SOMMAIRE|SUMMARY OF|CANADA REVENUE|AGENCE DU REVENU"
    r"|REMUNERATION|PENSION|DISPOSITIONS|TUITION|PROTECTED B|PROT[EÉ]G[EÉ] B",
    re.IGNORECASE,
)


def extract_form_code(text: str) -> str | None:
    """Recognized Canadian tax slip code (T4, T4A, T5008, ...), if present."""
    match = _FORM_CODE_RE.search(text.upper())
    if not match:
        return None
    if not _FORM_CONTEXT_RE.search(text):
        return None
    return match.group(1)


def extract_tax_year(text: str) -> str | None:
    """Tax year from a government slip (e.g. 'Year 2023' on a T4), not a transaction date."""
    for pattern in (_TAX_YEAR_RE, _TAX_YEAR_TITLE_RE):
        match = pattern.search(text)
        if match:
            year = int(match.group(1))
            if 1990 <= year <= 2100:
                return str(year)
    return None


def extract_store_name(text: str) -> str | None:
    """Known retail chain name if it appears anywhere in the OCR text."""
    for pattern, name in _STORE_PATTERNS:
        if pattern.search(text):
            return name
    return None


def extract_org_name(text: str) -> str | None:
    """First plausible employer/payer name line, skipping known form boilerplate.

    Used as a naming fallback for government tax slips, where the "largest
    text near the top" heuristic in extract_vendor tends to pick up form
    titles or bilingual labels instead of the actual organization name.
    """
    for line in text.splitlines():
        stripped = line.strip()
        # Civic addresses ("153 NIAGARA DR", "4700 KEELE STREET") start with
        # a street number; organization name lines don't, so this cheaply
        # tells an employer/payer name apart from the address line next to it.
        if not stripped or stripped[0].isdigit():
            continue
        # Check boilerplate/digit-heaviness per *candidate*, not per line:
        # OCR on boxed forms often merges the org name with an adjacent
        # bilingual label on one physical line ("YORK UNIVERSITY Year 2023
        # Statement of Remuneration Paid"), and the all-caps-words pattern
        # naturally only matches the "YORK UNIVERSITY" portion of that line
        # since the rest is mixed-case.
        for match in _ORG_NAME_RE.finditer(stripped):
            candidate = match.group(0)
            if _FORM_BOILERPLATE_RE.search(candidate):
                continue
            # Reject serial numbers / barcodes ("JTA9500902-0301613-..."),
            # which pass the ALL-CAPS-words shape but are mostly digits.
            digits = sum(c.isdigit() for c in candidate)
            if digits > len(candidate) * 0.2:
                continue
            return candidate
    return None


@dataclass
class OcrResult:
    text: str
    vendor: str | None
    date: str | None  # normalized YYYY-MM-DD
    time: str | None  # normalized HH:MM (24h), e.g. "15:25"
    total: str | None  # e.g. "81.52"
    mean_confidence: float  # 0-100, avg Tesseract word confidence
    form_code: str | None = None  # e.g. "T4", "T5008", if a known tax slip
    tax_year: str | None = None  # e.g. "2023" from "Year 2023" on a T4
    org_name: str | None = None  # employer/payer on tax slips
    doc_title: str | None = None  # recognized printed title slug, e.g. "transaction_history"
    party: str | None = None  # account holder / SHORTNAME when labeled
    page_number: int | None = None  # from a printed "PAGE X OF Y"
    page_total: int | None = None
    line_items: list[tuple[str, str]] | None = None  # (description, "12.34")


def _normalize_date(match: re.Match, kind: str) -> str | None:
    try:
        if kind == "ymd":
            y, m, d = (int(g) for g in match.groups())
        elif kind == "mdy":
            m, d, y = (int(g) for g in match.groups())
        elif kind == "month_name_dmy":
            month_str, day_str, year_str = match.groups()
            m = MONTHS.get(month_str.lower())
            d, y = int(day_str), int(year_str)
            if m is None:
                return None
        else:
            return None
        # Upper bound is next year, not 2100: a scanned document can't be
        # dated decades in the future, and OCR misreads (2023 -> 2053)
        # otherwise produce confidently wrong filenames.
        if not (1 <= m <= 12 and 1 <= d <= 31 and 1990 <= y <= datetime.now().year + 1):
            return None
        return f"{y:04d}-{m:02d}-{d:02d}"
    except (ValueError, TypeError):
        return None


def extract_date(text: str) -> str | None:
    """First plausible date found in OCR text, normalized to YYYY-MM-DD."""
    text = _O_BEFORE_DIGIT_RE.sub("0", text)
    text = _O_AFTER_DIGIT_RE.sub("0", text)
    for pattern, kind in _DATE_PATTERNS:
        match = pattern.search(text)
        if match:
            normalized = _normalize_date(match, kind)
            if normalized:
                return normalized
    return None


def extract_time(text: str) -> str | None:
    """First HH:MM clock time found in OCR text, normalized to zero-padded 24h."""
    match = _TIME_RE.search(text)
    if not match:
        return None
    hour, minute = int(match.group(1)), int(match.group(2))
    return f"{hour:02d}:{minute:02d}"


def extract_total(text: str) -> str | None:
    """Dollar amount on a line containing "total" but not "subtotal"."""
    lines = text.splitlines()
    for i, line in enumerate(lines):
        if not _TOTAL_LINE_RE.search(line):
            continue
        money = _MONEY_RE.search(line)
        if not money and i + 1 < len(lines):
            money = _MONEY_RE.search(lines[i + 1])
        if money:
            return money.group(1)
    return None


def extract_line_items(text: str) -> list[tuple[str, str]]:
    """SKU lines as (description, amount_str), skipping totals/tender noise.

    Amounts stay strings (same convention as extract_total). Categorization /
    split booking lives in taxes `scripts/receipt_itemize.py`.
    """
    out: list[tuple[str, str]] = []
    for raw in text.splitlines():
        line = re.sub(r"\s+", " ", raw).strip(" -\t")
        if not line or len(line) < 3 or _ITEM_SKIP_RE.search(line):
            continue
        match = _ITEM_LINE_RE.match(line)
        if not match:
            continue
        desc = match.group(1).strip(" -$")
        if not desc or _ITEM_SKIP_RE.search(desc):
            continue
        if len(re.sub(r"[^A-Za-z]", "", desc)) < 2:
            continue
        out.append((re.sub(r"\s+", " ", desc), match.group(2)))
    return out


def extract_vendor(
    image_bgr: np.ndarray,
    lang: str = "eng",
    top_fraction: float = 0.4,
    min_word_confidence: float = 40.0,
    psm: int = 6,
) -> str | None:
    """Largest-font text line within the top portion of the crop.

    Receipt headers and letterheads are typically both topmost and printed
    larger than body text, so we group Tesseract's word boxes into lines and
    pick the tallest line above the confidence floor.
    """
    height = image_bgr.shape[0]
    cutoff = int(height * top_fraction)
    region = image_bgr[:cutoff] if cutoff > 0 else image_bgr

    config = f"--psm {psm}"
    data = pytesseract.image_to_data(
        region, lang=lang, config=config, output_type=pytesseract.Output.DICT
    )

    lines: dict[tuple[int, int, int], list[int]] = {}
    for i, text in enumerate(data["text"]):
        if not text.strip():
            continue
        try:
            conf = float(data["conf"][i])
        except (TypeError, ValueError):
            continue
        if conf < min_word_confidence:
            continue
        key = (data["block_num"][i], data["par_num"][i], data["line_num"][i])
        lines.setdefault(key, []).append(i)

    if not lines:
        return _first_nonblank_line(
            pytesseract.image_to_string(region, lang=lang, config=config)
        )

    def line_stats(indices: list[int]) -> tuple[float, float]:
        # Median, not max: a single descender/ascender glyph (comma, "y",
        # parenthesis) can inflate one word's bounding box well above the
        # line's actual font size, which would otherwise skew line selection.
        height = statistics.median(data["height"][i] for i in indices)
        top = min(data["top"][i] for i in indices)
        return height, top

    stats = {key: line_stats(idxs) for key, idxs in lines.items()}
    max_height = max(h for h, _ in stats.values())
    # Treat near-max heights as the same font size and break ties by
    # topmost line, rather than picking whichever line measured tallest.
    tolerance = max(3.0, max_height * 0.15)
    candidates = [key for key, (h, _) in stats.items() if h >= max_height - tolerance]
    best_key = min(candidates, key=lambda k: stats[k][1])

    words = [data["text"][i].strip() for i in sorted(lines[best_key])]
    vendor = " ".join(w for w in words if w)
    return vendor or None


def _first_nonblank_line(text: str) -> str | None:
    for line in text.splitlines():
        stripped = line.strip()
        if stripped:
            return stripped
    return None


def run_ocr(
    image_bgr: np.ndarray,
    lang: str = "eng",
    vendor_top_fraction: float = 0.4,
    min_word_confidence: float = 40.0,
    psm: int = 6,
) -> OcrResult:
    """Run Tesseract once for full text, once (on a crop) for vendor detection."""
    config = f"--psm {psm}"
    text = pytesseract.image_to_string(image_bgr, lang=lang, config=config)

    data = pytesseract.image_to_data(
        image_bgr, lang=lang, config=config, output_type=pytesseract.Output.DICT
    )
    confidences = [float(c) for c in data["conf"] if c not in ("-1", -1)]
    mean_confidence = sum(confidences) / len(confidences) if confidences else 0.0

    vendor = clean_vendor(
        extract_vendor(
            image_bgr,
            lang=lang,
            top_fraction=vendor_top_fraction,
            min_word_confidence=min_word_confidence,
            psm=psm,
        )
    )
    date = extract_date(text)
    time = extract_time(text)
    total = extract_total(text)
    tax_year = extract_tax_year(text)

    store = extract_store_name(text)
    if store and (not vendor or store.upper().split("-")[0].split()[0] not in vendor.upper()):
        vendor = store

    form_code = extract_form_code(text)
    org_name = extract_org_name(text) if form_code else None
    if form_code:
        vendor = f"{form_code} {org_name}" if org_name else form_code

    page_marker = extract_page_marker(text)
    line_items = extract_line_items(text) or None

    return OcrResult(
        text=text,
        vendor=vendor,
        date=date,
        time=time,
        total=total,
        mean_confidence=mean_confidence,
        form_code=form_code,
        tax_year=tax_year,
        org_name=org_name,
        doc_title=extract_doc_title(text),
        party=extract_party(text),
        page_number=page_marker[0] if page_marker else None,
        page_total=page_marker[1] if page_marker else None,
        line_items=line_items,
    )


def run_ocr_on_path(path: str, **kwargs) -> OcrResult:
    image = cv2.imread(path)
    if image is None:
        raise ValueError(f"Could not read image: {path}")
    return run_ocr(image, **kwargs)