"""Build the final training/eval datasets.

Stage 1 — repair: replace each accepted truncated abstract with its
recovered text, normalized to the corpus format (the corpus strips all
punctuation except periods; sentence-final periods are free-standing ' . '
tokens; intra-token periods like '94.5' survive). The transform is
validated against every accepted pair by comparing the normalized recovered
prefix to the corpus text we already hold.

Stage 2 — titles: prepend the paper title, `<title> | <abstract>`, applied
identically to every split.

Usage: python build_datasets.py
Inputs: data/raw/*, data/external/{recovered_*,pii_title_cache.jsonl,*_with_pii.csv}
Outputs: data/recovered/{train,train_expanded,val,test}[_titled].csv
"""

import json
import re
import unicodedata
from difflib import SequenceMatcher
from pathlib import Path

import pandas as pd

REPO = Path(__file__).resolve().parent.parent
EXT = REPO / "data/external"
OUT = REPO / "data/recovered"

DASHES = "‐‑‒–—―−"
SECTION_HEADERS = re.compile(
    r"\b(BACKGROUND|METHODS?|RESULTS?|CONCLUSIONS?|OBJECTIVES?|PURPOSE"
    r"|INTRODUCTION|RATIONALE|AIMS?|DISCUSSION|MATERIALS|SIGNIFICANCE"
    r"|FINDINGS|INTERPRETATION|SETTING|DESIGN|PARTICIPANTS|MEASUREMENTS"
    r"|LIMITATIONS|IMPLICATIONS|HYPOTHESIS|UNLABELLED|IMPORTANCE"
    r"|EXPOSURES?|OUTCOMES?)(\s+AND\s+[A-Z]{4,})?\b[:.]?")


def strip_to_corpus_format(text):
    t = unicodedata.normalize("NFKC", str(text))
    t = SECTION_HEADERS.sub("", t)
    t = t.replace("!", ".").replace("?", ".")
    t = re.sub(r"(?<=\S)\.(?=\S)", "\x00", t)
    t = t.replace(".", " . ").replace("\x00", ".")
    t = t.replace("-", " ")
    for d in DASHES:
        t = t.replace(d, "")
    t = re.sub(r"[^A-Za-z0-9. ]", "", t)
    return " ".join(t.split())


def validate_transform():
    scores = []
    for split in ("train", "val", "test"):
        rec = pd.read_csv(EXT / f"recovered_{split}.csv")
        for _, r in rec[rec["accepted"]].iterrows():
            a = str(r["Abstract"]).split()[:-2]
            b = strip_to_corpus_format(r["RecoveredAbstract"]).split()[: len(a) + 20]
            m = SequenceMatcher(None, a, b, autojunk=False)
            scores.append(sum(bl.size for bl in m.get_matching_blocks()) / max(len(a), 1))
    mean = sum(scores) / len(scores)
    print(f"transform validation: {len(scores)} pairs, mean token agreement {mean:.4f}")
    assert mean >= 0.90, "transform does not reproduce corpus format"


def load_titles():
    titles = {}
    for line in (EXT / "pii_title_cache.jsonl").read_text().splitlines():
        try:
            r = json.loads(line)
            if r.get("title"):
                titles[r["pii"]] = r["title"]
        except (json.JSONDecodeError, KeyError):
            continue
    return titles


def clean_title(raw):
    t = strip_to_corpus_format(raw)
    while t.endswith(" ."):
        t = t[:-2].rstrip()
    return t


def main():
    validate_transform()
    OUT.mkdir(exist_ok=True)
    titles = load_titles()

    repl = {}
    for split in ("train", "val", "test"):
        rec = pd.read_csv(EXT / f"recovered_{split}.csv")
        acc = rec[rec["accepted"]]
        for pii, a in zip(acc["PII"], acc["RecoveredAbstract"]):
            repl[pii] = strip_to_corpus_format(a)
    extra = EXT / "recovered_expanded_extra.csv"
    if extra.exists():
        rec = pd.read_csv(extra)
        acc = rec[rec["accepted"]]
        for pii, a in zip(acc["PII"], acc["RecoveredAbstract"]):
            repl[pii] = strip_to_corpus_format(a)

    pii_by_split = {
        s: dict(zip(*(lambda d: (d["Filename"], d["PII"]))(
            pd.read_csv(EXT / f"{s}_with_pii.csv"))))
        for s in ("train", "val", "test")
    }

    for name in ("train", "val", "test", "train_expanded"):
        df = pd.read_csv(REPO / f"data/raw/{name}.csv")
        piis = (df["PaperID"] if name == "train_expanded"
                else df["Filename"].map(pii_by_split[name.replace("_expanded", "")]))
        df["Abstract"] = [repl.get(p, a) for p, a in zip(piis, df["Abstract"])]
        df.to_csv(OUT / f"{name}.csv", index=False)

        titled = []
        for pii, ab in zip(piis, df["Abstract"]):
            t = titles.get(pii)
            titled.append(f"{clean_title(t)} | {ab}" if t else str(ab))
        df_t = df.copy()
        df_t["Abstract"] = titled
        df_t.to_csv(OUT / f"{name}_titled.csv", index=False)

        trunc = (~df["Abstract"].astype(str).str.rstrip().str.endswith(".")).mean()
        print(f"{name}: {len(df)} rows, residual truncation {trunc:.1%}, "
              f"titled variant written")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())