scripts/build_dataset.py

"""Build a structured salary knowledge base from the IDA and Djøf 2025
lønstatistik PDFs.

Both PDFs are digital-born with reliable text layers, so we extract exact
digits via deterministic word-position parsing (extract_engine), not vision OCR.

Output: one JSON record per (table-row x statistic-bearing context), normalized
to a unified RAG-friendly schema (see schema.json). We keep the percentile
fields the user asked for (p25 / median / p75) plus mean, p90, count, and rich
provenance + a generated natural-language `rag_text` per record.
"""
from __future__ import annotations

import json
import re
import sys
from pathlib import Path

import pdfplumber

sys.path.insert(0, str(Path(__file__).parent))
from extract_engine import (  # noqa: E402
    cluster_rows, dedupe_glyphs, extract_lattice_table, extract_page_table,
    get_rulings, norm,
)

RAW = Path("/home/simonl/Desktop/competiton/pay-equity-for-eu/data/raw")
OUT = Path("/home/simonl/Desktop/competiton/pay-equity-for-eu/data/processed/lonstatistik")

IDA_PDF = RAW / "ida-loenstatistik-2025.pdf"
DJOF_PDF = RAW / "Privatansatte lnstatistik 2025.pdf"

CURRENT_YEAR = 2025

# ---------------------------------------------------------------------------
# table context (caption) extraction
# ---------------------------------------------------------------------------

SECTOR_HINTS = [
    (("staten", "ministerie"), "public_state"),
    (("kommun",), "public_municipal"),
    (("region",), "public_regional"),
    (("offentlig",), "public"),
    (("selvstænd", "selvstaend"), "self_employed"),
    (("privat",), "private"),
]


def guess_sector(text: str) -> str:
    t = norm(text)
    for keys, sector in SECTOR_HINTS:
        if any(k in t for k in keys):
            return sector
    return "private"  # both booklets are private-employee focused by default


PERIOD_RE = re.compile(r"(januar|februar|marts|april|maj|juni|juli|august|"
                       r"september|oktober|november|december)\s*20\d{2}", re.I)


def guess_period(text: str) -> str | None:
    m = PERIOD_RE.search(text)
    return m.group(0) if m else None


# breakdown dimension from "fordelt på X" / "Opdelt efter X"
DIM_RE = re.compile(r"(?:fordelt på|opdelt efter|opdelt i|ift\.?|i forhold til)\s+([^.]+)", re.I)


def guess_dimension(text: str) -> str | None:
    m = DIM_RE.search(text)
    if not m:
        return None
    d = re.sub(r"\s+", " ", m.group(1)).strip()
    return d[:80]


TABEL_RE = re.compile(r"TABEL\s*(\d+)\s*:?\s*", re.I)


def dehyphenate(text: str) -> str:
    """Join words split across line breaks: 'stil- lingstype' -> 'stillingstype'."""
    return re.sub(r"(\w)[\u00ad-]\s+(\w)", r"\1\2", text)


def ida_page_context(pg, rulings):
    """IDA caption lives in the left margin (x < first ruling)."""
    words = pg.extract_words(keep_blank_chars=False)
    cap = [w for w in words if w["x1"] <= rulings[0] + 1]
    cap.sort(key=lambda w: (round(w["top"]), w["x0"]))
    text = re.sub(r"\s+", " ", " ".join(w["text"] for w in cap)).strip()
    text = dehyphenate(text)
    return parse_caption(text)


FORTSAT_RE = re.compile(r"\(?\s*fortsat\s*\)?", re.I)
HDR_TOKENS = ("antal", "gennemsnit", "median", "kvartil", "fraktil", "nedre", "ovre")


def _is_header_line(line: str) -> bool:
    n = norm(line)
    return sum(k in n for k in HDR_TOKENS) >= 2


def djof_page_context(pg, prev_ctx=None):
    """Djøf caption has a fixed top structure for the FIRST page of a table:
       row0: TABEL N            row1: title ('Månedsløn ift. branche')
       row2: period             row3: 'Opdelt efter DIMENSION'
       row4: table-level group/subgroup (e.g. 'HD', 'Jurister', 'Alle kandidater')

    Continuation pages ('TABEL N (FORTSAT)') inherit title/period/dimension from
    prev_ctx and only carry a subgroup line like 'Chefer (fortsat)'.
    """
    rows = cluster_rows(pg.extract_words(keep_blank_chars=False))
    lines = [re.sub(r"\s+", " ", " ".join(dedupe_glyphs(w["text"]) for w in r)).strip()
             for r in rows[:6]]
    head = " ".join(lines[:1])
    is_fortsat = bool(FORTSAT_RE.search(head)) and "TABEL" in head.upper()

    ctx = {
        "table_ids": [], "table_title": None, "sector": "private",
        "period": None, "dimension": None, "caption": " ".join(lines)[:400],
        "table_group": None, "is_continuation": is_fortsat,
    }
    m = TABEL_RE.search(lines[0] if lines else "")
    if m:
        ctx["table_ids"] = [int(m.group(1))]

    if is_fortsat and prev_ctx and prev_ctx.get("table_ids") == ctx["table_ids"]:
        # inherit from the table's first page
        ctx["table_title"] = prev_ctx.get("table_title")
        ctx["period"] = prev_ctx.get("period")
        ctx["dimension"] = prev_ctx.get("dimension")
        # subgroup line: typically '<subgroup> (fortsat) Antal Gennemsnit ...'
        # i.e. the header row carries a left-side subgroup prefix. Extract the
        # text before the first stat-header token.
        for ln in lines[1:]:
            if not ln or "tabel" in norm(ln):
                continue
            if norm(ln) in ("stilling bruttolon", "bruttolon", "stilling"):
                continue
            # cut at first header token to isolate the subgroup prefix
            sub = re.split(r"(?i)\b(Antal|Gennemsnit|Nedre|Median|Øvre|90%)\b", ln)[0]
            sub = FORTSAT_RE.sub("", sub).strip(" -()")
            sub = re.sub(r"^Stilling\s+", "", sub, flags=re.I).strip()
            if sub and not re.search(r"\d{3}", sub):
                ctx["table_group"] = clean_label(sub)
                break
        return ctx

    # first page of the table
    if len(lines) > 1:
        ctx["table_title"] = lines[1]
    if len(lines) > 2:
        ctx["period"] = guess_period(lines[2]) or (lines[2] if "20" in lines[2] else None)
    if len(lines) > 3:
        md = DIM_RE.search(lines[3])
        ctx["dimension"] = (re.sub(r"\s+", " ", md.group(1)).strip()[:60] if md
                            else (lines[3] if lines[3] else None))
    if len(lines) > 4:
        cand = lines[4]
        if cand and not _is_header_line(cand) and not any(
                k in norm(cand) for k in ("bruttolon", "stilling")):
            ctx["table_group"] = clean_label(cand)
    return ctx


def parse_caption(text: str) -> dict:
    table_ids = TABEL_RE.findall(text)
    # title: text right after first TABEL N: up to first period or 'Månedsløn'
    title = None
    m = TABEL_RE.search(text)
    if m:
        rest = text[m.end():]
        rest = re.split(r"(Månedsløn|Bruttoløn fordelt|Brutoløn|Grundløn|September 20|Maj 20|Opdelt|\.)", rest)[0]
        title = re.sub(r"\s+", " ", rest).strip(" .:-")
    # Only classify sector from the *defining* part of the caption (before the
    # first long explanatory sentence). Use text up to the period statement.
    head = text
    pm = PERIOD_RE.search(text)
    if pm:
        head = text[: pm.end() + 40]  # title + 'fordelt på X' + period (+ a little)
    return {
        "table_ids": [int(t) for t in table_ids],
        "table_title": title or None,
        "sector": guess_sector(head),
        "period": guess_period(text),
        "dimension": guess_dimension(text),
        "caption": text[:400],
    }


# ---------------------------------------------------------------------------
# row-label parsing -> experience / age / segment
# ---------------------------------------------------------------------------

YEAR_RANGE_RE = re.compile(r"^(?:før\s*)?(\d{4})(?:\s*[-–]\s*(\d{2,4}))?$", re.I)
AGE_RE = re.compile(r"(\d{2})\s*[-–]\s*(\d{2})\s*år|(\d{2})\s*år", re.I)


def _clean(label: str) -> str:
    """Light clean: soft-hyphen -> hyphen, collapse spaces (keeps Danish chars)."""
    return re.sub(r"\s+", " ", label.replace("\u00ad", "-")).strip()


def parse_grad_year(label: str):
    """Return (grad_year_start, grad_year_end) from a cohort label like
    '1986-1987', 'Før 1986', '2024', '2012-14'."""
    l = norm(label).replace(" ", "")
    before = l.startswith("for")  # 'før'
    m = re.search(r"(\d{4})(?:[-–](\d{2,4}))?", l)
    if not m:
        return None, None, before
    y0 = int(m.group(1))
    y1 = m.group(2)
    if y1:
        y1 = int(y1) if len(y1) == 4 else int(str(y0)[:2] + y1)
    else:
        y1 = y0
    return y0, y1, before


def experience_from_grad(y0, y1, before):
    """Years of experience from graduation cohort (proxy: years since graduation)."""
    if y0 is None:
        return None, None
    if before:
        # 'Før 1986' -> graduated before 1986 -> >= ~39 yrs experience
        return CURRENT_YEAR - y0, None  # open-ended lower bound
    exp_max = CURRENT_YEAR - y0
    exp_min = CURRENT_YEAR - y1
    return exp_min, exp_max


def parse_age(label: str):
    m = AGE_RE.search(_clean(label))
    if not m:
        return None, None
    if m.group(1):
        return int(m.group(1)), int(m.group(2))
    if m.group(3):
        return int(m.group(3)), None  # '60 år og derover' -> open-ended
    return None, None


# label classification: is this a cohort row, age row, or a category row?
def label_kind(label: str, dimension: str | None):
    c = _clean(label)
    lc = c.lower()
    # age must be checked before year (both contain digits)
    if re.search(r"\bår\b", lc) and re.search(r"\d", lc):
        return "age"
    if (re.search(r"\b(19|20)\d{2}\b", c) or lc.startswith("før ")
            or YEAR_RANGE_RE.match(c.replace(" ", ""))):
        return "graduation_year"
    if lc in ("alle", "i alt", "ialt"):
        return "total"
    return "category"


# ---------------------------------------------------------------------------
# normalization
# ---------------------------------------------------------------------------

STAT_FIELDS = ["count", "mean", "p25", "median", "p75", "p90",
               "base", "supplement", "pension", "bonus_avg", "bonus_pct"]


def clean_label(s: str) -> str:
    s = s.replace("\xad", "-")
    # join words split across line breaks: 'program- mering' -> 'programmering'
    s = re.sub(r"(\w)-\s+(\w)", r"\1\2", s)
    s = re.sub(r"\s+", " ", s).strip()
    # strip trailing/leading stray dashes (missing-value placeholders that
    # leaked into a label) and surrounding punctuation
    s = re.sub(r"(\s*-\s*)+$", "", s).strip()
    s = re.sub(r"^(\s*-\s*)+", "", s).strip()
    return s


# canonical panel segment labels for the dual-panel Djøf tables
_PANEL_CANON = [
    ("kandidat", "Kandidater (master's degree)"),
    ("mellemudd", "Mid-level education (HA, HD, BA)"),
    ("direktør", "Directors"),
    ("ikke-chef", "Non-managers"),
    ("chef", "Managers"),
]


def clean_panel_title(title: str) -> str | None:
    """Collapse a noisy panel title to a canonical segment label."""
    if not title:
        return None
    t = norm(title)
    for key, canon in _PANEL_CANON:
        if key in t:
            return canon
    return clean_label(title)


def build_records_ida(pg, pidx, ctx, table):
    cols = table["columns"]
    col_stat = {i: c["stat"] for i, c in enumerate(cols)}
    col_hdr = {i: c.get("header", "") for i, c in enumerate(cols)}
    # detect if this is a "matrix" table: data columns are categories (no stat)
    stat_cols = [i for i, s in col_stat.items() if s]
    matrix = len(stat_cols) <= 1 and len(cols) > 3
    records = []
    dim = ctx.get("dimension")
    for row in table["rows"]:
        label = clean_label(row["label"])
        kind = label_kind(label, dim)
        recs_for_row = []
        if matrix:
            # each numeric cell is a gross-monthly value for category=col header
            for ci, val in row["cells"].items():
                seg = clean_label(col_hdr.get(ci, f"col{ci}"))
                if not seg or norm(seg) in ("alle",):
                    seg_label = seg or None
                else:
                    seg_label = seg
                rec = base_record(ctx, pidx, "IDA")
                fill_label_dims(rec, label, kind)
                rec["measure"] = {"mean": val}
                rec["segment_dimension"] = dim or "stillingstype"
                rec["segment_label"] = seg_label
                recs_for_row.append(rec)
        else:
            measure = {}
            for ci, val in row["cells"].items():
                s = col_stat.get(ci)
                if s:
                    measure[s] = val
            if not measure:
                continue
            rec = base_record(ctx, pidx, "IDA")
            fill_label_dims(rec, label, kind)
            rec["measure"] = measure
            rec["segment_dimension"] = dim
            rec["segment_label"] = label if kind == "category" else None
            recs_for_row.append(rec)
        records.extend(recs_for_row)
    return records


def build_records_djof(pg, pidx, ctx, table):
    cols = table["columns"]
    panels = table.get("panels", [{"col_range": [0, len(cols)], "title": ""}])
    records = []
    dim = ctx.get("dimension")
    for row in table["rows"]:
        label = clean_label(row["label"])
        kind = label_kind(label, dim)
        for p in panels:
            s, e = p["col_range"]
            measure = {}
            for ci, val in row["cells"].items():
                if s <= ci < e:
                    st = cols[ci]["stat"]
                    if st:
                        measure[st] = val
            if not measure:
                continue
            rec = base_record(ctx, pidx, "Djøf")
            fill_label_dims(rec, label, kind)
            rec["measure"] = measure
            rec["table_group"] = ctx.get("table_group")
            seg = clean_panel_title(p.get("title", ""))
            # if multiple panels, panel title is the segment; else the row label
            if len(panels) > 1:
                rec["segment_dimension"] = "education_or_position_group"
                rec["segment_label"] = seg
            else:
                rec["segment_dimension"] = dim if kind == "category" else None
                rec["segment_label"] = label if kind == "category" else None
            records.append(rec)
    return records


def base_record(ctx, pidx, union):
    return {
        "union": union,
        "source_doc": "IDA Lønstatistik 2025" if union == "IDA" else "Djøf Privatansatte Lønstatistik 2025",
        "source_page": pidx + 1,
        "table_id": ctx.get("table_ids", [None])[0] if ctx.get("table_ids") else None,
        "table_title": ctx.get("table_title"),
        "sector": ctx.get("sector"),
        "period": ctx.get("period"),
        "currency": "DKK",
        "pay_concept": None,
        "row_dimension": ctx.get("dimension"),
        "row_label": None,
        "graduation_year_start": None,
        "graduation_year_end": None,
        "experience_years_min": None,
        "experience_years_max": None,
        "age_min": None,
        "age_max": None,
        "table_group": ctx.get("table_group"),
        "segment_dimension": None,
        "segment_label": None,
        "measure": {},
    }


def fill_label_dims(rec, label, kind):
    rec["row_label"] = label
    if kind == "graduation_year":
        y0, y1, before = parse_grad_year(label)
        rec["graduation_year_start"] = y0
        rec["graduation_year_end"] = None if before else y1
        emin, emax = experience_from_grad(y0, y1, before)
        rec["experience_years_min"] = emin
        rec["experience_years_max"] = emax
    elif kind == "age":
        a0, a1 = parse_age(label)
        rec["age_min"] = a0
        rec["age_max"] = a1


NON_SALARY_TITLE_KW = ("lønudvikling", "udvikling", "fremskriv", "forventning",
                        "resultat", "omsætning", "omsaetning", "ejerandel",
                        "feriedage", "arbejdstid", "stigning")


def is_salary_record(rec):
    """Heuristic gate: keep only records that represent actual salary levels.

    Drops percentage/growth tables, projections, and business-metric tables
    (self-employed turnover/result), and rows whose monetary values are
    implausible as salaries.
    """
    title = norm(rec.get("table_title") or "")
    if any(k in title for k in NON_SALARY_TITLE_KW):
        # keep self-employed *income* tables (årsindkomst) though
        if "indkomst" not in title and "årsind" not in title and "arsind" not in title:
            return False
    m = rec["measure"]
    money = [v for k, v in m.items()
             if k in ("mean", "median", "p25", "p75", "p90", "base")]
    if not money:
        return False
    # corrupted multi-panel matrix rows: an implausibly large respondent count
    # on a non-total row indicates adjacent-panel numbers were merged.
    rl = norm(rec.get("row_label") or "")
    if rl not in ("alle", "i alt", "ialt") and m.get("count", 0) > 20000:
        return False
    pc = rec.get("pay_concept")
    if pc == "annual_income":
        return all(1000 <= v <= 100_000_000 for v in money)
    # monthly: plausible Danish gross monthly salary band
    return all(5000 <= v <= 500_000 for v in money)


def infer_pay_concept(rec):
    cap = norm((rec.get("table_title") or "") + " " + (rec.get("row_dimension") or ""))
    m = rec["measure"]
    if "base" in m or "supplement" in m:
        return "gross_monthly_with_components"
    if "selvstænd" in norm(rec.get("sector") or "") or rec.get("sector") == "self_employed":
        return "annual_income"
    if "nettolon" in cap:
        return "net_monthly"
    return "gross_monthly"


PERCENTILE_NAMES = {"p25": "25th percentile", "median": "median",
                    "p75": "75th percentile", "p90": "90th percentile",
                    "mean": "average", "base": "base salary",
                    "supplement": "supplement", "pension": "pension contribution",
                    "bonus_avg": "average annual bonus", "bonus_pct": "share receiving bonus"}


def make_rag_text(rec):
    m = rec["measure"]
    union = rec["union"]
    parts = [f"{union} 2025 salary statistics"]
    if rec.get("table_title"):
        parts.append(f"({rec['table_title'].title()})")
    seg = []
    if rec.get("segment_label"):
        seg.append(rec["segment_label"])
    if rec.get("row_label") and rec["row_label"] != rec.get("segment_label"):
        seg.append(rec["row_label"])
    sector = {"private": "private sector", "public_state": "state sector",
              "public_municipal": "municipal sector", "public_regional": "regional sector",
              "public": "public sector", "self_employed": "self-employed"}.get(rec.get("sector"), rec.get("sector"))
    head = f"{parts[0]} {parts[1] if len(parts)>1 else ''}".strip()
    body = f"For {', '.join(seg) if seg else 'all members'} in the {sector}"
    if rec.get("experience_years_min") is not None:
        if rec.get("experience_years_max") is not None:
            body += f" (~{rec['experience_years_min']}-{rec['experience_years_max']} yrs experience)"
        else:
            body += f" (~{rec['experience_years_min']}+ yrs experience)"
    elif rec.get("age_min") is not None:
        body += f" (age {rec['age_min']}{'-'+str(rec['age_max']) if rec.get('age_max') and rec['age_max']!=rec['age_min'] else '+'})"
    pieces = []
    cur = rec["currency"]
    pay = rec.get("pay_concept", "gross_monthly")
    unit = "DKK/month" if "monthly" in (pay or "") else ("DKK/year" if pay == "annual_income" else "DKK")
    for k in ["mean", "p25", "median", "p75", "p90"]:
        if k in m:
            pieces.append(f"{PERCENTILE_NAMES[k]} {m[k]:,} {unit}")
    extra = []
    if "base" in m:
        extra.append(f"base {m['base']:,}")
    if "supplement" in m:
        extra.append(f"supplement {m['supplement']:,}")
    if "pension" in m:
        extra.append(f"pension {m['pension']:,}")
    if "bonus_pct" in m:
        extra.append(f"{m['bonus_pct']}% receive bonus")
    if "bonus_avg" in m:
        extra.append(f"avg bonus {m['bonus_avg']:,}/yr")
    n = f" Based on {m['count']} respondents." if "count" in m else ""
    sal = "; ".join(pieces)
    ex = (" " + ", ".join(extra) + ".") if extra else ""
    return f"{head}. {body}: {sal}.{ex}{n} Source: {rec['source_doc']}, p.{rec['source_page']}.".replace("  ", " ")


def process_pdf(path, union):
    records = []
    prev_ctx = None
    with pdfplumber.open(path) as pdf:
        for pidx, pg in enumerate(pdf.pages):
            rulings = get_rulings(pg)
            if union == "IDA" and rulings:
                ctx = ida_page_context(pg, rulings)
                if ctx.get("table_ids"):
                    prev_ctx = ctx
                table = extract_lattice_table(pg, rulings)
                if table and table["rows"]:
                    records += build_records_ida(pg, pidx, ctx, table)
            else:
                if union == "Djøf":
                    ctx = djof_page_context(pg, prev_ctx)
                    # remember the first-page context for continuation pages
                    if not ctx.get("is_continuation") and ctx.get("table_ids"):
                        prev_ctx = ctx
                else:
                    ctx = ida_page_context(pg, rulings or [40])
                table = extract_page_table(pg)
                if table and table["rows"]:
                    records += build_records_djof(pg, pidx, ctx, table)
    # finalize: pay_concept, salary gate, rag_text + ids
    kept = []
    dropped = 0
    for rec in records:
        rec["pay_concept"] = infer_pay_concept(rec)
        if not is_salary_record(rec):
            dropped += 1
            continue
        kept.append(rec)
    for i, rec in enumerate(kept):
        rec["rag_text"] = make_rag_text(rec)
        rec["id"] = f"{union.lower().replace('ø','o')}-p{rec['source_page']:03d}-{i:04d}"
    print(f"  {union}: dropped {dropped} non-salary/implausible records")
    return kept


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    all_recs = []
    for path, union in [(IDA_PDF, "IDA"), (DJOF_PDF, "Djøf")]:
        recs = process_pdf(path, union)
        print(f"{union}: {len(recs)} records")
        all_recs += recs
    (OUT / "salary_records.json").write_text(
        json.dumps(all_recs, ensure_ascii=False, indent=2))
    with (OUT / "salary_records.jsonl").open("w") as f:
        for r in all_recs:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"TOTAL: {len(all_recs)} records -> {OUT}")


if __name__ == "__main__":
    main()