scripts/extract_engine.py

"""Deterministic salary-table extraction engine for Danish lønstatistik PDFs.

Both IDA (lattice) and Djøf (borderless) are digital-born PDFs with reliable
text layers, so we extract exact digits by word-position clustering rather than
vision OCR (which risks digit errors in a salary knowledge base).

Pipeline per page:
  1. cluster words into rows by y (top)
  2. detect numeric data columns by clustering x-centers of numeric tokens
  3. map header words above each column to a canonical statistic key
  4. emit one record per (row-label x column) numeric cell? No -> one record per
     data row, with a dict of {stat_key: value}
"""
from __future__ import annotations

import re
import unicodedata
from dataclasses import dataclass, field

import pdfplumber

# ---- numeric parsing -------------------------------------------------------
# Danish format: thousands '.', decimal ','  e.g. 81.629  5,0%  45.171
NUM_RE = re.compile(r"^-?\d{1,3}(?:\.\d{3})*(?:,\d+)?%?$|^-?\d+(?:,\d+)?%?$")
DASH = {"-", "\u2013", "\u2014", "\u2212", "."}  # placeholders for missing


def is_number(tok: str) -> bool:
    t = tok.strip()
    if t in DASH:
        return True  # missing-value marker
    return bool(NUM_RE.match(t))


def parse_number(tok: str):
    """Return (value, is_pct, is_missing)."""
    t = tok.strip()
    if t in DASH or t == "":
        return None, False, True
    pct = t.endswith("%")
    t = t.rstrip("%")
    t = t.replace(".", "").replace(",", ".")
    try:
        v = float(t)
    except ValueError:
        return None, pct, True
    if v.is_integer():
        v = int(v)
    return v, pct, False


# ---- row clustering --------------------------------------------------------
def cluster_rows(words, ytol=3.0):
    """Group words into rows by their vertical center."""
    ws = sorted(words, key=lambda w: (round(w["top"]), w["x0"]))
    rows = []
    cur = []
    cur_y = None
    for w in ws:
        yc = (w["top"] + w["bottom"]) / 2
        if cur_y is None or abs(yc - cur_y) <= ytol:
            cur.append(w)
            ys = [(x["top"] + x["bottom"]) / 2 for x in cur]
            cur_y = sum(ys) / len(ys)
        else:
            rows.append(cur)
            cur = [w]
            cur_y = yc
    if cur:
        rows.append(cur)
    for r in rows:
        r.sort(key=lambda w: w["x0"])
    return rows


def norm(s: str) -> str:
    s = s.lower().strip()
    # fold Danish letters explicitly (NFKD does not decompose ø)
    s = (s.replace("ø", "o").replace("æ", "ae").replace("å", "aa")
           .replace("\xad", ""))
    s = unicodedata.normalize("NFKD", s)
    s = "".join(c for c in s if not unicodedata.combining(c))
    return s


# ---- header -> canonical stat key ------------------------------------------
def classify_stat(header_text: str) -> str | None:
    h = norm(header_text).replace("-", "").replace(" ", "")
    # order matters
    if "antal" in h:
        return "count"
    if "procent" in h and "bonus" in h:
        return "bonus_pct"
    if ("far" in h and "bonus" in h):
        return "bonus_pct"
    if "arligbonus" in h or "bonus" in h:
        return "bonus_avg"
    if "basislon" in h or "grundlon" in h:
        return "base"
    if "tillaeg" in h or "tillg" in h:
        return "supplement"
    if "pension" in h:
        return "pension"
    if "90" in h:
        return "p90"
    if "75" in h or "ovrekvartil" in h or h == "ovre":
        return "p75"
    if "median" in h or "50" in h:
        return "median"
    if "25" in h or "nedrekvartil" in h or h == "nedre":
        return "p25"
    if "gennemsnit" in h or "bruttolon" in h or "brutolon" in h or "nettolon" in h:
        return "mean"
    return None


def xcenter(w):
    return (w["x0"] + w["x1"]) / 2


def row_numeric_tokens(row):
    return [w for w in row if is_number(w["text"])]


def cluster_columns(centers, gap=14.0):
    """1-D clustering of x-centers into columns. Returns list of (lo,hi,mid)."""
    if not centers:
        return []
    centers = sorted(centers)
    cols = [[centers[0]]]
    for c in centers[1:]:
        if c - cols[-1][-1] <= gap:
            cols[-1].append(c)
        else:
            cols.append([c])
    return [(min(c), max(c), sum(c) / len(c)) for c in cols]


def _overlap(a0, a1, b0, b1):
    return max(0.0, min(a1, b1) - max(a0, b0))


def get_rulings(pg, min_count=4):
    """Return sorted x-positions of vertical ruling lines, or [] if borderless."""
    vx = sorted({round(e["x0"]) for e in pg.edges if e["orientation"] == "v"})
    # merge near-duplicates
    merged = []
    for x in vx:
        if not merged or x - merged[-1] > 3:
            merged.append(x)
    return merged if len(merged) >= min_count else []


def extract_lattice_table(pg, rulings, ytol=3.0):
    """Extract a ruled table (IDA): columns = gaps between vertical rulings.

    Column 0 (first ruling..second) is the row label; rulings[0] is the left
    edge of the table content, so caption text at x<rulings[0] is excluded.
    """
    left = rulings[0]
    right = rulings[-1]
    bounds = rulings  # column i spans bounds[i]..bounds[i+1]
    ncol = len(bounds) - 1
    if ncol < 3:
        return None

    words = [w for w in pg.extract_words(keep_blank_chars=False)
             if w["x0"] >= left - 2 and w["x1"] <= right + 2]
    rows = cluster_rows(words, ytol=ytol)

    def colof(w):
        cx = xcenter(w)
        for i in range(ncol):
            if bounds[i] - 1 <= cx <= bounds[i + 1] + 1:
                return i
        return None

    # classify data rows (>=3 numeric) vs header rows.
    # A row is a HEADER (not data) if it contains stat keywords, even if it has
    # numeric tokens like "25%"/"50%" (these are column titles, not values).
    HDR_KW = ("fraktil", "kvartil", "antal", "median", "gennemsnit",
              "bruttoløn", "basisløn", "grundløn", "tillæg", "pension",
              "bonus", "procent", "stilling", "årgang", "alder", "branche",
              "region", "løntrin", "stillingsniveau")

    def is_header_row(r):
        txt = norm(" ".join(w["text"] for w in r))
        return any(k in txt for k in HDR_KW)

    data_rows, header_rows = [], []
    for r in rows:
        nums = [w for w in r if is_number(w["text"])]
        if len(nums) >= 3 and not is_header_row(r):
            data_rows.append(r)
        else:
            header_rows.append(r)
    if len(data_rows) < 2:
        return None

    first_data_top = min(min(w["top"] for w in r) for r in data_rows)
    # build header text per column from ALL words above first data row
    # (captures multi-line category headers in matrix tables).
    col_headers = [[] for _ in range(ncol)]
    for w in words:
        # include numeric tokens too: "25%"/"50%" are part of headers
        if w["top"] < first_data_top - 1:
            ci = colof(w)
            if ci is not None and ci > 0:  # col 0 is label
                col_headers[ci].append((w["top"], w["x0"], w["text"]))
    col_stats, col_labels = [], []
    for ci in range(ncol):
        hs = sorted(col_headers[ci])
        # strip page-header noise (all-caps section words, page numbers)
        toks = [t for _, _, t in hs if t not in ("PRIVATANSATTE", "OFFENTLIGT",
                "ANSATTE", "SELVSTÆNDIG", "GENERELT")]
        htext = " ".join(toks)
        col_stats.append(classify_stat(htext))
        col_labels.append(re.sub(r"\s+", " ", htext).strip())

    out_rows = []
    for r in data_rows:
        cells = {i: [] for i in range(ncol)}
        for w in r:
            ci = colof(w)
            if ci is not None:
                cells[ci].append(w)
        # label = col 0 text
        label = re.sub(r"\s+", " ", " ".join(
            w["text"] for w in sorted(cells[0], key=lambda w: w["x0"]))).strip()
        cellvals, raw = {}, {}
        for ci in range(1, ncol):
            toks = sorted(cells[ci], key=lambda w: w["x0"])
            if not toks:
                continue
            txt = "".join(t["text"] for t in toks)
            v, pct, missing = parse_number(txt)
            if missing:
                continue
            cellvals[ci] = v
            raw[ci] = txt
        if label and cellvals:
            out_rows.append({"label": label, "cells": cellvals, "raw": raw})

    return {
        "columns": [{"bounds": [bounds[i], bounds[i + 1]], "stat": col_stats[i],
                     "header": col_labels[i]} for i in range(ncol)],
        "rows": out_rows,
        "mode": "lattice",
    }


STAT_HEADER_TOKENS = {
    "antal": "count", "gennemsnit": "mean", "median": "median",
    "nedre": "p25", "ovre": "p75", "90%-fraktil": "p90", "90%fraktil": "p90",
    "90%": "p90", "bruttoløn": "mean",
}


def _token_stat(text):
    t = norm(text).replace(" ", "")
    if "antal" in t:
        return "count"
    if "gennemsnit" in t or "bruttolon" in t:
        return "mean"
    if t == "nedre":
        return "p25"
    if "median" in t:
        return "median"
    if t == "ovre":
        return "p75"
    if "90" in t and "fraktil" in t:
        return "p90"
    return None


def find_header_anchors(rows):
    """Find the statistics header band in a borderless Djøf table.

    Scans header rows for stat tokens (Antal, Gennemsnit, Nedre, Median, Øvre,
    90%-fraktil) which may be spread across 2-3 stacked rows, and merges them
    into one ordered list of anchors. Handles dual-panel pages (two side-by-side
    blocks) by simply returning all anchors ordered by x.

    Returns (last_header_row_index, anchors) where anchors is a list of
    (stat_key, x_center, label).
    """
    # locate the band: rows containing >=2 stat tokens
    hit_rows = []
    for ri, r in enumerate(rows):
        n = sum(1 for w in r if _token_stat(w["text"]))
        if n >= 2:
            hit_rows.append(ri)
    if not hit_rows:
        return None, None
    lo, hi = min(hit_rows), max(hit_rows)
    # the band spans lo..hi (typically 1-3 rows); collect all stat tokens
    anchors = []
    for ri in range(lo, hi + 1):
        for w in rows[ri]:
            s = _token_stat(w["text"])
            if s:
                anchors.append((s, xcenter(w), w["text"]))
    anchors.sort(key=lambda a: a[1])
    return hi, anchors


def dedupe_glyphs(s):
    """Fix pdfplumber doubled-glyph artifact: 'BBrruuttttoolløønn' -> 'Bruttoløn'.

    Only collapses when a token is entirely pairs of identical chars.
    """
    out = []
    for tok in s.split():
        if len(tok) >= 4 and len(tok) % 2 == 0 and all(
                tok[i] == tok[i + 1] for i in range(0, len(tok), 2)):
            out.append(tok[::2])
        else:
            out.append(tok)
    return " ".join(out)


def split_panels(anchors):
    """Split anchors into panels at large x-gaps where the stat sequence restarts.

    A new panel starts when we see 'count' or 'mean' again after already having
    seen a full-ish block, or when there's a big x-gap. Returns list of
    (start_idx, end_idx) anchor index ranges.
    """
    if not anchors:
        return []
    panels = []
    start = 0
    seen = set()
    for i, (stat, x, _) in enumerate(anchors):
        if i > start and stat in ("count", "mean") and stat in seen:
            panels.append((start, i))
            start = i
            seen = set()
        seen.add(stat)
    panels.append((start, len(anchors)))
    return panels


_TITLE_SKIP = {"nedre", "ovre", "kvartil", "antal", "gennemsnit", "median",
               "bruttolon", "for", "fraktil"}


def find_panel_titles(rows, hri, first_anchor_x):
    """Find panel/segment titles (e.g. 'Direktører', 'kandidater') above headers.

    Picks fragments from the 1-2 rows directly above the stat band that sit over
    the numeric region (x >= first_anchor_x - 40) and are not stat/structural
    words. Returns list of (x_center, title_fragment).
    """
    titles = []
    for ri in range(max(0, hri - 3), hri):
        for w in rows[ri]:
            if xcenter(w) < first_anchor_x - 40:
                continue
            t = dedupe_glyphs(w["text"])
            tn = norm(t).replace(".", "").replace("-", "")
            if tn in _TITLE_SKIP or is_number(t):
                continue
            titles.append((xcenter(w), t, ri))
    return titles


def extract_page_table(pg, ytol=3.0):
    """Extract a borderless statistics table (Djøf) by anchoring on the header band.

    Handles dual-panel pages by splitting anchors into panels at stat restarts.
    Each numeric data cell is assigned to the nearest header anchor by x-center.
    Missing values (dashes) are simply absent, so this is robust to sparse rows.
    """
    words = pg.extract_words(keep_blank_chars=False)
    if not words:
        return None
    rows = cluster_rows(words, ytol=ytol)
    hri, anchors = find_header_anchors(rows)
    if not anchors or len(anchors) < 3:
        return None
    anchor_x = [a[1] for a in anchors]
    anchor_stat = [a[0] for a in anchors]
    ncol = len(anchors)

    panels = split_panels(anchors)
    title_frags = find_panel_titles(rows, hri, anchor_x[0])

    def assign_col(w):
        return min(range(ncol), key=lambda i: abs(anchor_x[i] - xcenter(w)))

    # panel x-boundaries
    panel_bounds = []
    for (s, e) in panels:
        lo = anchor_x[s] - 30
        hi = anchor_x[e - 1] + 30
        panel_bounds.append((lo, hi))

    def panel_title(lo, hi):
        frags = [t for x, t, ri in title_frags if lo - 20 <= x <= hi + 20]
        return re.sub(r"\s+", " ", " ".join(frags)).strip()

    out_rows = []
    for ri, r in enumerate(rows):
        if ri <= hri:
            continue
        nums = [w for w in r if is_number(w["text"])]
        if len(nums) < 2:
            continue
        first_num_x = min(xcenter(w) for w in nums)
        label_words = [w for w in r if xcenter(w) < first_num_x and not is_number(w["text"])]
        label = re.sub(r"\s+", " ", " ".join(
            w["text"] for w in sorted(label_words, key=lambda w: w["x0"]))).strip()
        if not label:
            continue
        cellvals, raw = {}, {}
        for w in nums:
            v, pct, missing = parse_number(w["text"])
            if missing:
                continue
            ci = assign_col(w)
            cellvals[ci] = v
            raw[ci] = w["text"]
        if label and cellvals:
            out_rows.append({"label": label, "cells": cellvals, "raw": raw})

    if len(out_rows) < 2:
        return None
    return {
        "columns": [
            {"xcenter": round(anchor_x[i], 1), "stat": anchor_stat[i],
             "header": dedupe_glyphs(anchors[i][2])}
            for i in range(ncol)
        ],
        "panels": [
            {"col_range": [s, e], "x_range": [round(panel_bounds[pi][0], 1),
             round(panel_bounds[pi][1], 1)], "title": panel_title(*panel_bounds[pi])}
            for pi, (s, e) in enumerate(panels)
        ],
        "rows": out_rows,
        "mode": "borderless",
    }


if __name__ == "__main__":
    import sys, json
    fn = sys.argv[1]
    pi = int(sys.argv[2]) - 1
    mode = sys.argv[3] if len(sys.argv) > 3 else "table"
    with pdfplumber.open(fn) as pdf:
        pg = pdf.pages[pi]
        if mode == "raw":
            for r in cluster_rows(pg.extract_words(keep_blank_chars=False)):
                print("  ".join(f"{round(w['x0'])}:{w['text']}" for w in r))
        else:
            rul = get_rulings(pg)
            if rul:
                res = extract_lattice_table(pg, rul)
            else:
                res = extract_page_table(pg)
            print(json.dumps(res, ensure_ascii=False, indent=2))