fund_rfid_data/build_rdf_dataset.py

#!/usr/bin/env python3
"""
SEC Fund Prospectus -> RDF Triple dataset builder (proof of concept).

Produces text->triples samples where:
  - INPUT  = the prospectus prose for a fund family (long natural-language text),
  - OUTPUT = a graph of entity->entity RDF triples (NOT flat literal attributes),
  - GOLD   = N-CEN structured filings (service providers) + the prospectus structure.

The graph edges are genuine relationships between entities:
    Fund   seriesOf       Trust
    Fund   advisedBy      InvestmentAdviser
    Fund   subAdvisedBy   SubAdviser
    Fund   transferAgent  TransferAgent
    Fund   custodian      Custodian
    Fund   administrator  Administrator
    Trust  underwrittenBy Distributor
    Fund   tracksIndex    Index           (flag-derived; index name from prose)

Gold for advisedBy/transferAgent/custodian/administrator/underwrittenBy comes
directly from N-CEN (no model). The prospectus prose is fetched from EDGAR and
serves as the model input.

Stages:
  gold     -- build the per-fund gold graph from the local N-CEN flat files
  fetch    -- download prospectus prose (485BPOS) for the selected N-CEN registrants
  samples  -- join prose + gold into text->triple training samples
  all      -- run gold, fetch, samples in order

Usage:
  python build_rdf_dataset.py gold     --ncen data/ncen/2025q3
  python build_rdf_dataset.py fetch    --limit 25
  python build_rdf_dataset.py samples
  python build_rdf_dataset.py all      --limit 25
"""

import argparse
import csv
import gzip
import json
import logging
import re
import sys
import time
import urllib.request
from collections import defaultdict
from pathlib import Path

log = logging.getLogger("rdf")

HEADERS = {"User-Agent": "FundDataResearch/1.0 research@university.edu",
           "Accept-Encoding": "gzip, deflate"}
EDGAR_RATE = 0.15  # seconds between SEC requests (well under 10/s limit)
_last = 0.0

OUT_DIR = Path("data/rdf_poc")
GOLD_PATH = OUT_DIR / "gold_graphs.jsonl"
PROSE_DIR = OUT_DIR / "prose"
SAMPLES_PATH = OUT_DIR / "samples.jsonl"

# Prose shorter than this is a stub: a 497 supplement fragment, or a closed-end
# fund (which files N-2/486BPOS, not the 485BPOS our fetcher targets) where only
# a tiny amendment was retrieved. Such text cannot ground a fund's triples, so
# the sample is skipped rather than emitted with non-extractable targets.
MIN_PROSE_CHARS = 50000

# RDF marker tokens from the thesis (Section 5.2)
T_START, T_PRED, T_OBJ, T_END = (
    "<triple_start>", "<predicate_marker>", "<object_marker>", "<triple_end>")


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def _get(url: str, timeout: int = 120) -> bytes:
    global _last
    dt = time.time() - _last
    if dt < EDGAR_RATE:
        time.sleep(EDGAR_RATE - dt)
    req = urllib.request.Request(url, headers=HEADERS)
    resp = urllib.request.urlopen(req, timeout=timeout)
    data = resp.read()
    _last = time.time()
    if resp.headers.get("Content-Encoding") == "gzip" or data[:2] == b"\x1f\x8b":
        data = gzip.decompress(data)
    return data


def _read_tsv(path: Path) -> list[dict]:
    if not path.exists():
        log.warning("missing %s", path)
        return []
    # csv.field_size_limit because some SEC fields are huge
    csv.field_size_limit(min(sys.maxsize, 2**31 - 1))
    with open(path, "r", encoding="utf-8", errors="replace", newline="") as f:
        return list(csv.DictReader(f, delimiter="\t"))


def _slug(s: str) -> str:
    """Make an IRI-safe local name from an entity string."""
    s = re.sub(r"[^A-Za-z0-9]+", "_", (s or "").strip())
    return s.strip("_") or "x"


def html_to_text(raw: str) -> str:
    raw = re.sub(r"(?is)<script.*?</script>", " ", raw)
    raw = re.sub(r"(?is)<style.*?</style>", " ", raw)
    txt = re.sub(r"(?s)<[^>]+>", " ", raw)
    txt = re.sub(r"&#\d+;", " ", txt)
    txt = re.sub(r"&[a-zA-Z]+;", " ", txt)
    txt = re.sub(r"[ \t]+", " ", txt)
    txt = re.sub(r"\s*\n\s*", "\n", txt)
    return re.sub(r"\n{3,}", "\n\n", txt).strip()


# --------------------------------------------------------------------------- #
# stage: gold  -- build per-registrant (trust) gold graph from N-CEN
# --------------------------------------------------------------------------- #
def build_gold(ncen_dir: Path, custodian_scope: str = "primary"):
    d = Path(ncen_dir)
    submission = {r["ACCESSION_NUMBER"]: r for r in _read_tsv(d / "SUBMISSION.tsv")}
    registrant = {r["ACCESSION_NUMBER"]: r for r in _read_tsv(d / "REGISTRANT.tsv")}

    # FUND_REPORTED_INFO: FUND_ID -> fund row (also gives ACCESSION + flags)
    funds = _read_tsv(d / "FUND_REPORTED_INFO.tsv")
    fund_by_id = {r["FUND_ID"]: r for r in funds}
    funds_by_acc = defaultdict(list)
    for r in funds:
        funds_by_acc[r["ACCESSION_NUMBER"]].append(r)

    # provider tables keyed by FUND_ID
    def by_fund(fname, name_col, lei_col):
        out = defaultdict(list)
        for r in _read_tsv(d / fname):
            fid = r.get("FUND_ID", "")
            nm = (r.get(name_col, "") or "").strip()
            if fid and nm:
                out[fid].append({"name": nm, "lei": (r.get(lei_col, "") or "").strip(),
                                 "affiliated": r.get("IS_AFFILIATED", ""),
                                 "type": r.get("ADVISER_TYPE", ""),
                                 "is_sub_custodian": r.get("IS_SUB_CUSTODIAN", "")})
        return out

    advisers = by_fund("ADVISER.tsv", "ADVISER_NAME", "ADVISER_LEI")
    tagents = by_fund("TRANSFER_AGENT.tsv", "TRANSFERAGENT_NAME", "TRANSFERAGENT_LEI")
    custodians = by_fund("CUSTODIAN.tsv", "CUSTODIAN_NAME", "CUSTODIAN_LEI")
    admins = by_fund("ADMIN.tsv", "ADMIN_NAME", "ADMIN_LEI")

    # Share classes per FUND_ID (the 3rd hierarchy level: Trust->Fund->ShareClass).
    # Each class has a full name and a ticker, both typically in the prospectus.
    share_classes = defaultdict(list)
    for r in _read_tsv(d / "SHARES_OUTSTANDING.tsv"):
        fid = r.get("FUND_ID", "")
        cid = (r.get("CLASS_ID", "") or "").strip()
        nm = (r.get("CLASS_NAME", "") or "").strip()
        if fid and cid:
            share_classes[fid].append(
                {"class_id": cid, "name": nm,
                 "ticker": (r.get("TICKER", "") or "").strip()})

    # Custodian scoping. Foreign sub-custodians (IS_SUB_CUSTODIAN=Y, ~88% of rows)
    # appear ONLY in N-CEN, never in the prospectus prose, and dominate the edge
    # count (66% of all edges). They are therefore unextractable noise for a
    # text->triples task. Default keeps the PRIMARY custodian only.
    #   "primary" -> only IS_SUB_CUSTODIAN != Y  (one prose-grounded edge per fund)
    #   "all"     -> every custodian row (legacy behaviour)
    #   "none"    -> drop the custodian relation entirely
    if custodian_scope == "primary":
        custodians = {fid: [c for c in cs
                            if (c.get("is_sub_custodian", "") or "").upper() != "Y"]
                      for fid, cs in custodians.items()}
    elif custodian_scope == "none":
        custodians = {}

    # underwriter (distributor) keyed by ACCESSION (trust level)
    underwriters = defaultdict(list)
    for r in _read_tsv(d / "PRINCIPAL_UNDERWRITER.tsv"):
        acc = r.get("ACCESSION_NUMBER", "")
        nm = (r.get("UNDERWRITER_NAME", "") or "").strip()
        if acc and nm:
            underwriters[acc].append({"name": nm,
                                      "lei": (r.get("UNDERWRITER_LEI", "") or "").strip()})

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    n_graphs = 0
    n_edges = 0
    with open(GOLD_PATH, "w", encoding="utf-8") as out:
        for acc, reg in registrant.items():
            cik = (reg.get("CIK", "") or "").strip().zfill(10)
            trust_name = (reg.get("REGISTRANT_NAME", "") or "").strip()
            if not trust_name:
                continue
            trust_iri = "trust:" + _slug(trust_name)

            triples = []
            entities = {trust_iri: {"type": "Trust", "label": trust_name,
                                    "lei": (reg.get("LEI", "") or "").strip()}}

            # trust-level distributor edges
            for u in underwriters.get(acc, []):
                d_iri = "org:" + _slug(u["name"])
                entities[d_iri] = {"type": "Distributor", "label": u["name"], "lei": u["lei"]}
                triples.append((trust_iri, "underwrittenBy", d_iri))

            for fr in funds_by_acc.get(acc, []):
                fid = fr["FUND_ID"]
                fname = (fr.get("FUND_NAME", "") or "").strip()
                if not fname:
                    continue
                f_iri = "fund:" + _slug(fname)
                entities[f_iri] = {"type": "Fund", "label": fname,
                                   "series_id": (fr.get("SERIES_ID", "") or "").strip(),
                                   "lei": (fr.get("LEI", "") or "").strip()}
                triples.append((f_iri, "seriesOf", trust_iri))

                if (fr.get("IS_INDEX", "") or "").upper() == "Y":
                    entities[f_iri]["is_index"] = True
                if (fr.get("IS_ETF", "") or "").upper() == "Y":
                    entities[f_iri]["is_etf"] = True

                for a in advisers.get(fid, []):
                    o = "org:" + _slug(a["name"])
                    is_sub = (a.get("type", "") or "").lower().startswith("sub")
                    entities[o] = {"type": "SubAdviser" if is_sub else "InvestmentAdviser",
                                   "label": a["name"], "lei": a["lei"]}
                    triples.append((f_iri, "subAdvisedBy" if is_sub else "advisedBy", o))
                for a in tagents.get(fid, []):
                    o = "org:" + _slug(a["name"])
                    entities[o] = {"type": "TransferAgent", "label": a["name"], "lei": a["lei"]}
                    triples.append((f_iri, "transferAgent", o))
                for a in custodians.get(fid, []):
                    o = "org:" + _slug(a["name"])
                    entities[o] = {"type": "Custodian", "label": a["name"], "lei": a["lei"]}
                    triples.append((f_iri, "custodian", o))
                for a in admins.get(fid, []):
                    o = "org:" + _slug(a["name"])
                    entities[o] = {"type": "Administrator", "label": a["name"], "lei": a["lei"]}
                    triples.append((f_iri, "administrator", o))

                # share classes (3rd level): Fund -hasShareClass-> Class,
                # Class -ticker-> "TICKER", Class -className-> "Full Name".
                # ticker/className objects are literals (stored as label).
                for sc in share_classes.get(fid, []):
                    c_iri = "class:" + sc["class_id"]
                    entities[c_iri] = {"type": "ShareClass",
                                       "label": sc["name"] or sc["class_id"],
                                       "class_id": sc["class_id"]}
                    triples.append((f_iri, "hasShareClass", c_iri))
                    if sc["name"]:
                        lit = "lit:" + _slug(sc["name"])
                        entities[lit] = {"type": "Literal", "label": sc["name"]}
                        triples.append((c_iri, "className", lit))
                    if sc["ticker"]:
                        lit = "lit:" + _slug(sc["ticker"])
                        entities[lit] = {"type": "Literal", "label": sc["ticker"]}
                        triples.append((c_iri, "ticker", lit))

            # dedupe triples
            triples = sorted(set(triples))
            if not triples:
                continue
            rec = {
                "accession": acc, "cik": cik,
                "trust_name": trust_name, "trust_iri": trust_iri,
                "n_funds": len(funds_by_acc.get(acc, [])),
                "entities": entities,
                "triples": [{"s": s, "p": p, "o": o} for s, p, o in triples],
            }
            out.write(json.dumps(rec, ensure_ascii=False) + "\n")
            n_graphs += 1
            n_edges += len(triples)

    log.info("gold: %d trust graphs, %d entity->entity edges -> %s",
             n_graphs, n_edges, GOLD_PATH)
    # quick edge-type histogram
    hist = defaultdict(int)
    with open(GOLD_PATH, encoding="utf-8") as f:
        for line in f:
            for t in json.loads(line)["triples"]:
                hist[t["p"]] += 1
    print("\nEdge-type histogram (entity->entity edges):")
    for p, c in sorted(hist.items(), key=lambda x: -x[1]):
        print(f"  {p:16s} {c:>7,}")
    print(f"\nTotal: {n_graphs:,} trust graphs, {n_edges:,} edges")


# --------------------------------------------------------------------------- #
# stage: fetch  -- prospectus prose for selected registrants
# --------------------------------------------------------------------------- #
# Full statutory prospectuses (cover ALL funds of a book) vs. short supplements.
FULL_PROSPECTUS_FORMS = ("485BPOS", "485APOS")
SUPPLEMENT_FORMS = ("497", "497K")
# A full prospectus book is large; post-effective AMENDMENTS share the 485BPOS
# form but are small (only the changed sections). Require at least this size to
# count a 485BPOS as a complete book rather than an amendment.
FULL_BOOK_MIN_BYTES = 2_000_000


def _prospectus_filings(cik: str, max_filings: int):
    """Return the current prospectus document set for a CIK.

    The gold (N-CEN) is CURRENT, so the source text must be too. A prospectus is a
    living document: the statutory book (485BPOS) plus running SUPPLEMENTS/addenda
    (497) that amend it (e.g. a changed custodian). We therefore return:
      1. the NEWEST *full* book -- the newest 485BPOS/485APOS above a size
         threshold (so a thin post-effective amendment is not mistaken for the
         book; older-but-complete books are skipped in favour of the current one);
      2. ALL 497/497K supplements filed AFTER that book's date -- the running
         changes a reader would apply on top of it.
    Concatenated, this is the up-to-date document a shareholder actually reads,
    aligned in time with the N-CEN gold. Falls back to the largest 485BPOS (any
    size) or, last, recent supplements, if no full book exists.
    Returns a list of (accession_nodash, primary_doc), book first.
    """
    url = f"https://data.sec.gov/submissions/CIK{cik}.json"
    j = json.loads(_get(url).decode("utf-8", "replace"))
    recent = j.get("filings", {}).get("recent", {})
    forms = recent.get("form", [])
    accs = recent.get("accessionNumber", [])
    docs = recent.get("primaryDocument", [])
    sizes = recent.get("size", [0] * len(forms))
    dates = recent.get("filingDate", [""] * len(forms))

    books = []      # (date, acc, doc, size) for full-form prospectuses
    supps = []      # (date, acc, doc) for supplements
    for i, fm in enumerate(forms):
        if not docs[i]:
            continue
        if fm in FULL_PROSPECTUS_FORMS:
            books.append((dates[i], accs[i].replace("-", ""), docs[i], sizes[i] or 0))
        elif fm in SUPPLEMENT_FORMS:
            supps.append((dates[i], accs[i].replace("-", ""), docs[i]))

    # newest FULL book (above the size threshold); else largest book of any size
    full_books = [b for b in books if b[3] >= FULL_BOOK_MIN_BYTES]
    if full_books:
        book = max(full_books, key=lambda b: b[0])           # newest complete book
    elif books:
        book = max(books, key=lambda b: b[3])                # largest available
    else:
        # no full prospectus at all -> most recent supplements only
        supps.sort(reverse=True)
        return [(a, d) for _dt, a, d in supps[:max_filings]]

    book_date = book[0]
    chosen = [(book[1], book[2])]
    # all supplements filed on/after the book's date, newest first
    later = sorted([s for s in supps if s[0] >= book_date], reverse=True)
    for _dt, a, d in later:
        if len(chosen) >= max_filings:
            break
        chosen.append((a, d))
    return chosen


# An SAI is identified by the appointment sentences it contains for the very
# service-provider roles that the summary prospectus omits.
_SAI_SIGNAL = re.compile(
    r"serves as (?:custodian|administrator|transfer agent|sub-administrator)"
    r"|as custodian to|as administrator to|statement of additional information",
    re.I)


def _sai_documents(cik: str, acc: str, primary_doc: str):
    """Return extra HTM documents in an accession that look like the SAI.

    The SAI is filed as a separate document inside the same 485BPOS accession as
    the prospectus. We list the accession's index, fetch each non-primary HTM
    document, and keep those whose text contains role-appointment sentences
    (\\code{serves as custodian}, etc.). Returns a list of extracted-text strings.
    """
    try:
        idx = json.loads(_get(
            f"https://www.sec.gov/Archives/edgar/data/{int(cik)}/{acc}/index.json"
        ).decode("utf-8", "replace"))
    except Exception:
        return []
    out = []
    items = idx.get("directory", {}).get("item", [])
    for it in items:
        name = it.get("name", "")
        if not name.lower().endswith((".htm", ".html")):
            continue
        if name == primary_doc or name.lower().startswith("r"):  # skip primary + XBRL R-files
            continue
        try:
            raw = _get(
                f"https://www.sec.gov/Archives/edgar/data/{int(cik)}/{acc}/{name}"
            ).decode("utf-8", "replace")
        except Exception:
            continue
        if not _SAI_SIGNAL.search(raw):
            continue
        txt = html_to_text(raw)
        if len(txt) >= 5000:
            out.append(txt)
    return out


def fetch_prose(limit: int, max_filings: int = 8, include_sai: bool = True,
                only_ciks=None):
    """Fetch ALL recent prospectus filings per trust and concatenate per CIK.

    For each selected trust, up to `max_filings` recent prospectus filings are
    downloaded and their extracted text concatenated into data/rdf_poc/prose/
    <cik>.txt, separated by a form-feed marker. This maximises the chance that
    every N-CEN fund of the trust has its prospectus section present, raising
    per-fund segmentation coverage.

    When `include_sai` is set (default), the Statement of Additional Information
    (SAI) document(s) in each accession are also appended. The SAI names the
    administrator, transfer agent and (primary) custodian in role-appointment
    sentences that the summary prospectus omits, which is what makes those
    relations prose-grounded; foreign sub-custodians remain absent and so stay
    correctly excluded by context-grounding.
    """
    if not GOLD_PATH.exists():
        log.error("run `gold` first"); return
    PROSE_DIR.mkdir(parents=True, exist_ok=True)
    rows = [json.loads(l) for l in open(GOLD_PATH, encoding="utf-8")]
    # prefer trusts with the most edges (richer graphs) for the PoC slice
    rows.sort(key=lambda r: -len(r["triples"]))
    if only_ciks:
        want = set(only_ciks)
        rows = [r for r in rows if r["cik"] in want]
    done = 0
    tot_filings = 0
    for r in rows:
        if done >= limit:
            break
        cik = r["cik"]
        out = PROSE_DIR / f"{cik}.txt"
        if out.exists():
            done += 1
            continue
        try:
            filings = _prospectus_filings(cik, max_filings)
            if not filings:
                log.info("no prospectus filing for CIK %s (%s)", cik, r["trust_name"])
                continue
            parts = []
            for acc, doc in filings:
                try:
                    url = f"https://www.sec.gov/Archives/edgar/data/{int(cik)}/{acc}/{doc}"
                    raw = _get(url).decode("utf-8", "replace")
                    txt = html_to_text(raw)
                    if len(txt) >= 2000:
                        parts.append(txt)
                    if include_sai:
                        parts.extend(_sai_documents(cik, acc, doc))
                except Exception as e:
                    log.debug("  filing %s failed: %s", acc, e)
            if not parts:
                log.info("no usable prospectus text for CIK %s, skipping", cik)
                continue
            combined = ("\n\n\f\n\n").join(parts)  # form-feed separates books
            out.write_text(combined, encoding="utf-8")
            tot_filings += len(parts)
            log.info("[%d/%d] %s  %d filings  prose=%d chars  %s",
                     done + 1, limit, cik, len(parts), len(combined),
                     r["trust_name"][:40])
            done += 1
        except Exception as e:
            log.warning("fetch failed for CIK %s: %s", cik, e)
    log.info("fetched %d filings across %d trusts -> %s",
             tot_filings, done, PROSE_DIR)


# --------------------------------------------------------------------------- #
# stage: samples  -- join prose + gold into text->triple samples
# --------------------------------------------------------------------------- #
def serialize_triples(triples, entities) -> str:
    """Render triples in the thesis's <triple_start> marker format, grouped by subject.

    Marker form (Models 2/4 with grammar-terminal tokens in the vocabulary).
    """
    bys = defaultdict(list)
    for t in triples:
        bys[t["s"]].append((t["p"], t["o"]))
    chunks = []
    for s, pos in bys.items():
        s_label = entities.get(s, {}).get("label", s)
        body = [f"{T_START} {s_label}"]
        for p, o in pos:
            o_label = entities.get(o, {}).get("label", o)
            body.append(f"{T_PRED} {p} {T_OBJ} {o_label}")
        body.append(T_END)
        chunks.append(" ".join(body))
    return "\n".join(chunks)


def serialize_triples_plain(triples, entities) -> str:
    """Render the SAME triples in a plain Turtle-like form with NO special tokens.

    For Models 1/3 (decoder-only / encoder-decoder without added grammar-terminal
    tokens). Subjects are factored out and predicate-object lists are separated by
    ';' and ',' exactly as in Turtle, so the two serializations encode identical
    content and differ only in whether the delimiters are dedicated tokens.
    """
    bys = defaultdict(list)
    for t in triples:
        bys[t["s"]].append((t["p"], t["o"]))
    chunks = []
    for s, pos in bys.items():
        s_label = entities.get(s, {}).get("label", s)
        by_pred = defaultdict(list)
        for p, o in pos:
            by_pred[p].append(entities.get(o, {}).get("label", o))
        preds = []
        for p, objs in by_pred.items():
            preds.append(f"{p} " + " , ".join(objs))
        chunks.append(f"{s_label} " + " ; ".join(preds) + " .")
    return "\n".join(chunks)


_GND_SUFFIX = re.compile(
    r"\b(llc|l\.l\.c|inc|incorporated|corp|corporation|company|co|ltd|limited"
    r"|lp|l\.p|llp|n\.a|na|the|trust)\b", re.I)


def _gnorm(s: str) -> str:
    """Normalize a name for prose-grounding checks (lowercase, strip legal suffixes)."""
    s = _GND_SUFFIX.sub(" ", (s or "").lower())
    s = re.sub(r"[^a-z0-9]+", " ", s)
    return re.sub(r"\s+", " ", s).strip()


# Keywords that signal each relation is being asserted *about* the named entity.
# Context-grounding requires one of these near the object name, so that a name
# mentioned in the wrong role (e.g. an adviser firm that is ALSO the administrator,
# but only ever called "adviser" in the text) is not counted for that relation.
#
# For roles whose bare word is too common in boilerplate ("...including the
# custodian, and the transfer agent..."), we require an APPOINTMENT phrasing
# ("serves as custodian", "acts as administrator") rather than the bare word, so
# that a generic mention near an unrelated name cannot count. This is what makes
# the safeguard genuinely catch the custodian case.
_REL_KEYWORDS = {
    "advisedBy":     ["advis", "investment manager", "the manager"],
    "subAdvisedBy":  ["sub-advis", "subadvis", "sub advis", "sub-adviser"],
    "transferAgent": ["transfer agent"],
    "underwrittenBy": ["distributor", "underwriter"],
    "seriesOf":      ["series of", "a series", "each a series"],
    # appointment-phrased (role word alone is boilerplate):
    "administrator": ["as administrator", "as the administrator", "administrator to",
                      "administrator for", "administrator of the", "administration services"],
    "custodian":     ["as custodian", "as the custodian", "custodian to",
                      "custodian for", "custodian of the", "serves as custodian",
                      "acts as custodian"],
}
_CTX_WINDOW = 200  # chars on each side of the name occurrence


def _context_grounded(name_norm, predicate, text_low):
    """True iff the object name occurs near a keyword for its relation in text_low.

    Uses the most distinctive token of the name (longest >=5 chars) to locate
    candidate positions, then checks a +/- window for any relation keyword. This
    distinguishes "X serves as administrator" (valid) from a bare mention of X.
    """
    kws = _REL_KEYWORDS.get(predicate, [])
    if not kws or not name_norm:
        return False
    toks = [w for w in name_norm.split() if len(w) >= 5] or name_norm.split()
    if not toks:
        return False
    tok = toks[0]
    for m in re.finditer(re.escape(tok), text_low):
        win = text_low[max(0, m.start() - _CTX_WINDOW): m.start() + _CTX_WINDOW]
        if any(k in win for k in kws):
            return True
    return False


# Stopwords for fuzzy name grounding (generic words shared by many firm names).
_GND_STOP = {"of", "and", "the", "fund", "funds", "services", "service", "asset",
             "management", "investment", "investments", "company", "co", "group",
             "capital", "advisors", "advisers", "global", "trust", "inc", "llc"}


def _name_tokens(name_norm):
    return [w for w in name_norm.split() if w not in _GND_STOP and len(w) > 2]


def _fuzzy_grounded(name_norm, ntext_tokens):
    """True iff a strong majority of the name's distinctive tokens occur in text.

    This is the robust middle ground between exact-substring (which false-positives
    on shared brand words like "John Hancock") and the strict appointment-keyword
    check (which false-negatives on valid paraphrases). A name counts as grounded
    when >=2/3 of its non-stopword tokens appear in the document, AND at least one
    distinctive (long) token does -- so a fund whose administrator is literally
    named is captured, while an unrelated brand mention is not.
    """
    toks = _name_tokens(name_norm)
    if not toks:
        return False
    present = sum(1 for w in toks if w in ntext_tokens)
    if present / len(toks) < 0.67:
        return False
    # require at least one *distinctive* token (length >= 5) to be present,
    # so two short generic tokens cannot ground a name on their own
    return any(len(w) >= 5 and w in ntext_tokens for w in toks)


_ALIAS_STOP = {"the", "inc", "llc", "lp", "llp", "company", "co", "corp",
               "corporation", "trust", "fund", "funds", "of", "and", "group",
               "na", "national", "association", "incorporated", "limited", "ltd"}


def _alias_tokens(label):
    return [w for w in re.sub(r"[^a-z0-9 ]", " ", (label or "").lower()).split()
            if w not in _ALIAS_STOP and len(w) >= 2]


def _acronym(label):
    return "".join(w[0] for w in _alias_tokens(label))


def _alias_grounded(label, text_norm):
    """True iff the entity NAME occurs in the text, tolerant to surface variants.

    This is exactly the dataset requirement: the triple's value (the entity name)
    must appear in the document, allowing variation (legal-form suffixes, filler
    words, common abbreviations) but not mere scattered token overlap. Two ways to
    match, both requiring the name to appear *coherently*:
      1. the significant tokens occur consecutively (allowing up to two filler
         words between them), e.g. "State Street Bank and Trust Company" matches
         "State Street ... Trust";
      2. the acronym of the significant tokens (>=3 letters) occurs as a word,
         e.g. "BNY"/"JPM".
    Validated against a 122B LLM verifier: where they disagree, the name was in
    fact present verbatim, i.e. this rule is the faithful test of the requirement.
    `text_norm` is the lowercased text with punctuation collapsed to spaces.
    """
    tk = _alias_tokens(label)
    if not tk:
        return False
    # Cheap gate first: the most distinctive token must appear at all. Substring
    # `in` is far faster than regex over multi-MB texts, so this avoids running the
    # flexible pattern on the vast majority of (name, document) pairs.
    distinctive = max(tk, key=len)
    if (" " + distinctive) not in (" " + text_norm):
        ac = _acronym(label)
        return len(ac) >= 3 and (" " + ac + " ") in (" " + text_norm + " ")
    if len(tk) == 1:
        return re.search(r"\b" + re.escape(tk[0]) + r"\b", text_norm) is not None
    pat = r"\b" + r"\W+(?:\w+\W+){0,2}?".join(re.escape(w) for w in tk) + r"\b"
    if re.search(pat, text_norm):
        return True
    ac = _acronym(label)
    return len(ac) >= 3 and re.search(r"\b" + re.escape(ac) + r"\b", text_norm) is not None


def annotate_grounding(triples, entities, text):
    """Annotate each triple with grounding flags against `text`.

    A triple is only a valid TEXT->triple target if its object NAME appears in the
    input. Flags recorded per triple:
      - `alias_grounded`   : the name occurs coherently, tolerant to legal-form
                             suffixes / filler words / acronyms (the recommended,
                             requirement-faithful filter; fast, deterministic);
      - `grounded`         : the normalized name occurs verbatim (strict substring);
      - `context_grounded` : the name occurs near an appointment keyword for its
                             relation (strictest; e.g. "X serves as custodian").
    Returns (annotated_triples, n_alias, n_grounded, n_context).
    """
    ntext = _gnorm(text)
    text_norm = re.sub(r"[^a-z0-9 ]", " ", text.lower())
    tlow = text.lower()
    out = []
    n_alias = n_grounded = n_ctx = 0
    for t in triples:
        olbl = entities.get(t["o"], {}).get("label", "")
        no = _gnorm(olbl)
        al = _alias_grounded(olbl, text_norm)
        g = bool(no) and no in ntext
        cg = g and _context_grounded(no, t["p"], tlow)
        n_alias += int(al)
        n_grounded += int(g)
        n_ctx += int(cg)
        out.append({**t, "alias_grounded": al, "grounded": g, "context_grounded": cg})
    return out, n_alias, n_grounded, n_ctx


# How strict the build is about triples actually appearing in the input text.
# The build always ANNOTATES all flags; this only controls which (if any) is used
# to FILTER the standard dataset (the full, unfiltered set is always written
# alongside, see _build_samples_per_fund).
#   "alias" (default) -> keep triples whose object name occurs coherently in the
#                        text, tolerant to legal-form/filler/acronym variation.
#                        Fast, deterministic, and the faithful test of the
#                        requirement (validated against a 122B LLM verifier).
#   "name"            -> keep triples whose normalized name occurs verbatim.
#   "context"         -> strictest: name near an appointment keyword.
#   "none"            -> keep all triples (annotate only, no filtering).
#   "llm"             -> use the LLM match verdict (llm_extract.py match mode):
#                        keeps a triple only if a strong model confirmed the
#                        entity plays that ROLE in the text. Most accurate (catches
#                        right-name/wrong-role cases, e.g. a bank named as
#                        securities-lending agent but not custodian); requires a
#                        prior match run (see LLM_MATCH_PATH).
GROUNDING_MODE = "alias"

# Per-(cik,predicate,object-IRI) llm_grounded verdicts, loaded on demand from a
# match run. Keyed (cik, predicate, o_iri) -> bool.
LLM_MATCH_PATH = OUT_DIR / "match_all.jsonl"
_LLM_VERDICTS = None


def _load_llm_verdicts():
    global _LLM_VERDICTS
    if _LLM_VERDICTS is not None:
        return _LLM_VERDICTS
    _LLM_VERDICTS = {}
    if LLM_MATCH_PATH.exists():
        for line in open(LLM_MATCH_PATH, encoding="utf-8"):
            r = json.loads(line)
            cik = r.get("cik")
            for t in r.get("triples", []):
                iri = t.get("o_iri", t.get("o"))
                _LLM_VERDICTS[(cik, t.get("p"), iri)] = bool(t.get("llm_grounded"))
    log.info("loaded %d LLM match verdicts from %s", len(_LLM_VERDICTS), LLM_MATCH_PATH)
    return _LLM_VERDICTS


def filter_grounded(triples, cik=None):
    """Drop triples that fail the active GROUNDING_MODE. Returns kept triples.

    This is what *enforces* — rather than merely measures — that every target
    fact is present in the input, so a non-prose relation (like the foreign
    sub-custodians) cannot silently enter the standard dataset.
    """
    if GROUNDING_MODE == "llm":
        v = _load_llm_verdicts()
        return [t for t in triples if v.get((cik, t["p"], t["o"]), False)]
    if GROUNDING_MODE == "none":
        return triples
    key = {"context": "context_grounded", "name": "grounded",
           "alias": "alias_grounded"}[GROUNDING_MODE]
    return [t for t in triples if t.get(key)]


def ontology_schema(triples, entities) -> dict:
    """Inferred meta-schema (subject type -> predicate -> object type), per thesis 5.3."""
    schema = defaultdict(lambda: defaultdict(set))
    for t in triples:
        st = entities.get(t["s"], {}).get("type", "Thing")
        ot = entities.get(t["o"], {}).get("type", "Thing")
        schema[st][t["p"]].add(ot)
    return {st: {p: sorted(os) for p, os in preds.items()}
            for st, preds in schema.items()}


# --------------------------------------------------------------------------- #
# per-fund segmentation
# --------------------------------------------------------------------------- #
def _name_variants(fund_name: str):
    """Generate prospectus-heading variants of an N-CEN fund name.

    N-CEN names and prospectus headings often differ in the legal-form suffix
    (Fund vs ETF vs Portfolio) and in spacing/punctuation, so we match on a set
    of normalized variants rather than the exact string.
    """
    base = fund_name.strip()
    stems = {base}
    # swap the trailing legal-form word
    for suf in (" Fund", " ETF", " Portfolio", " Trust"):
        if base.endswith(suf):
            root = base[: -len(suf)]
            for alt in (" Fund", " ETF", " Portfolio", ""):
                stems.add(root + alt)
            stems.add(root)
            break
    # also the bare root with no suffix
    return {s.strip() for s in stems if len(s.strip()) >= 6}


# A fund section opens with the fund name immediately followed by a summary
# heading. Filers use several styles right after the name, so we accept any of:
#   - "Fund Summary" / "Summary" (mutual funds)
#   - "Investment Objective" / "Principal Investment Strateg" (objective heading)
#   - "The Fund seeks ..." (ETF objective sentence)
#   - "Class/Ticker:" or a ticker block / "(formerly ...)" (ETF & multi-class
#     summary headers, e.g. JPMorgan, Invesco ETF)
# The MIN_SEGMENT_CHARS guard below is the real safety net: any anchor that
# resolves to a collapsed (too-short) segment is discarded, so a slightly broad
# anchor set cannot reintroduce the spurious-cross-reference problem.
_SECTION_ANCHOR = re.compile(
    r"\s{0,5}("
    r"Fund Summary|Summary Section|Investment Objective|Principal Investment Strateg"
    r"|The Fund seeks|Class/Ticker|Ticker(s)?:|\(formerly"
    r")", re.I)
MIN_SEGMENT_CHARS = 1500  # a real fund summary is at least this long


def _heading_positions(text: str, fund_names):
    """Return all (offset, fund_name) anchored heading candidates in the text.

    A candidate is a position where a fund name variant is immediately followed
    by a strong section anchor. A fund may have several candidates (its book may
    be concatenated more than once, or it appears in multiple books); all are
    kept so segmentation can choose the one that yields a non-trivial segment.
    """
    cands = []
    for fn in fund_names:
        for v in sorted(_name_variants(fn), key=len, reverse=True):
            for m in re.finditer(re.escape(v), text):
                if _SECTION_ANCHOR.match(text[m.end():m.end() + 45]):
                    cands.append((m.start(), fn))
    return sorted(cands)


def _segment_prose(text: str, fund_names):
    """Split prose into per-fund segments using anchored section headings.

    Algorithm:
      1. collect all anchored heading candidates across the (possibly multi-book)
         text and sort them by offset;
      2. each candidate's segment runs to the NEXT candidate heading (of any
         fund), so boundaries come only from real section starts;
      3. for each fund, choose the candidate whose segment is longest and at
         least MIN_SEGMENT_CHARS, discarding collapsed (too-short) candidates.
    Returns {fund_name: segment_text} for funds with a usable section.
    """
    cands = _heading_positions(text, fund_names)
    if not cands:
        return {}
    offsets = [c[0] for c in cands]
    best = {}  # fund -> (length, segment_text)
    for i, (off, fn) in enumerate(cands):
        end = offsets[i + 1] if i + 1 < len(offsets) else len(text)
        seg = text[off:end]
        if len(seg) < MIN_SEGMENT_CHARS:
            continue
        if fn not in best or len(seg) > best[fn][0]:
            best[fn] = (len(seg), seg)
    return {fn: seg for fn, (ln, seg) in best.items()}


def build_samples(per_fund: bool = True):
    if not GOLD_PATH.exists():
        log.error("run `gold` first"); return
    if per_fund:
        return _build_samples_per_fund()
    gold = {json.loads(l)["cik"]: json.loads(l)
            for l in open(GOLD_PATH, encoding="utf-8")}
    n = n_skip = 0
    n_tot_triples = n_kept_triples = 0
    tot_text = tot_json = 0
    outfull = open(SAMPLES_FULL_PATH, "w", encoding="utf-8")
    with open(SAMPLES_PATH, "w", encoding="utf-8") as out:
        for prose_file in sorted(PROSE_DIR.glob("*.txt")):
            cik = prose_file.stem
            g = gold.get(cik)
            if not g:
                continue
            text = prose_file.read_text(encoding="utf-8")
            if len(text) < MIN_PROSE_CHARS:
                n_skip += 1
                continue   # stub prose (supplement / closed-end fund) — skip
            ents = g["entities"]
            cand, n_al, n_gnd, n_ctx = annotate_grounding(g["triples"], ents, text)
            n_tot_triples += len(cand)
            _write_sample(out, outfull, f"{cik}:ALL", cik, g["trust_name"],
                          None, "", False, text, cand, ents)
            kept = filter_grounded(cand, cik)
            if not kept:
                continue
            n_kept_triples += len(kept)
            target = serialize_triples(kept, ents)
            n += 1
            tot_text += len(text)
            tot_json += len(target)
    outfull.close()
    log.info("wrote %d samples -> %s (full set -> %s)", n, SAMPLES_PATH, SAMPLES_FULL_PATH)
    if n:
        agree = 100 * n_kept_triples / max(1, n_tot_triples)
        print(f"\n{n} trust samples ({n_skip} stub-prose skipped); grounding={GROUNDING_MODE}")
        print(f"  text<->gold agreement: {n_kept_triples}/{n_tot_triples} "
              f"triples = {agree:.0f}% (kept after grounding filter)")
        print(f"  mean input  : {tot_text//n:>8,} chars")
        print(f"  mean target : {tot_json//n:>8,} chars")
        print(f"  mean ratio  : {tot_text/max(1,tot_json):>8.1f} : 1  (text : json)")


SAMPLES_FULL_PATH = OUT_DIR / "samples_full.jsonl"


def _write_sample(out, outfull, sid, cik, trust_name, fund, series_id,
                  segmented, text, cand_triples, ents):
    """Write the filtered record to `out` and the full annotated record to `outfull`.

    `cand_triples` are all candidate triples already annotated with grounding
    flags. The standard file (`out`) carries only the grounding-filtered triples
    (every target is text-present); the full file (`outfull`) carries ALL triples
    with their flags, for traceability and alternative thresholds.
    """
    def ref_ents(triples):
        ref = set()
        for t in triples:
            ref.add(t["s"]); ref.add(t["o"])
        return {k: ents[k] for k in ref if k in ents}

    # full record (all triples + flags)
    fe = ref_ents(cand_triples)
    n_al = sum(t.get("alias_grounded") for t in cand_triples)
    n_g = sum(t.get("grounded") for t in cand_triples)
    n_cx = sum(t.get("context_grounded") for t in cand_triples)
    outfull.write(json.dumps({
        "sample_id": sid, "cik": cik, "trust_name": trust_name, "fund": fund,
        "series_id": series_id, "segmented": segmented, "input_text": text,
        "ontology": ontology_schema(cand_triples, fe),
        "target_triples": cand_triples,
        "stats": {"input_chars": len(text), "n_triples": len(cand_triples),
                  "n_alias_grounded": n_al, "n_grounded": n_g,
                  "n_context_grounded": n_cx},
    }, ensure_ascii=False) + "\n")

    # standard record (filtered triples only)
    kept = filter_grounded(cand_triples, cik)
    if not kept:
        return
    se = ref_ents(kept)
    target = serialize_triples(kept, se)
    rec = {
        "sample_id": sid, "cik": cik, "trust_name": trust_name, "fund": fund,
        "series_id": series_id, "segmented": segmented,
        "grounding_mode": GROUNDING_MODE,
        "input_text": text, "ontology": ontology_schema(kept, se),
        "target_triples": kept,
        "target_serialized": target,
        "target_serialized_plain": serialize_triples_plain(kept, se),
        "stats": {"input_chars": len(text), "target_chars": len(target),
                  "n_triples": len(kept), "n_entities": len(se),
                  "text_to_json_ratio": round(len(text) / max(1, len(target)), 1)},
    }
    out.write(json.dumps(rec, ensure_ascii=False) + "\n")


def _build_samples_per_fund():
    """One sample per fund: the fund's prospectus section -> the fund's subgraph.

    For each trust, the prose is segmented into per-fund sections. Each fund's
    target is its own edges (advisedBy, custodian, ...) plus the fund-anchored
    seriesOf edge and the trust-level underwrittenBy edge (shared, but a true
    fact about that fund's trust). Funds whose section cannot be located in the
    prose are skipped and counted; trusts where nothing can be located fall back
    to a single whole-trust sample so no data is silently dropped.

    Two files are written: SAMPLES_PATH (grounding-filtered, training-ready) and
    SAMPLES_FULL_PATH (all triples with grounding flags, for traceability).
    """
    gold = {json.loads(l)["cik"]: json.loads(l)
            for l in open(GOLD_PATH, encoding="utf-8")}
    n = 0
    n_funds_total = n_funds_located = n_fallback_trusts = n_dropped_empty = 0
    tot_text = tot_json = 0
    ratios = []
    outfull = open(SAMPLES_FULL_PATH, "w", encoding="utf-8")
    with open(SAMPLES_PATH, "w", encoding="utf-8") as out:
        for prose_file in sorted(PROSE_DIR.glob("*.txt")):
            cik = prose_file.stem
            g = gold.get(cik)
            if not g:
                continue
            text = prose_file.read_text(encoding="utf-8")
            if len(text) < MIN_PROSE_CHARS:
                continue   # stub prose (supplement / closed-end fund) — skip
            ents = g["entities"]
            # group triples by subject fund IRI; collect trust-level edges
            fund_iris = {e_iri for e_iri, e in ents.items() if e["type"] == "Fund"}
            trust_iri = g["trust_iri"]
            by_fund = defaultdict(list)
            trust_edges = []
            for t in g["triples"]:
                if t["s"] in fund_iris:
                    by_fund[t["s"]].append(t)
                elif t["s"] == trust_iri:
                    trust_edges.append(t)  # e.g. underwrittenBy

            fund_label = {iri: ents[iri]["label"] for iri in fund_iris}
            n_funds_total += len(fund_iris)
            segs = _segment_prose(text, list(fund_label.values()))
            label_to_iri = {v: k for k, v in fund_label.items()}

            if not segs:  # whole-trust fallback (no section located)
                cand = g["triples"]
                cand, n_al, n_gnd, n_ctx = annotate_grounding(cand, ents, text)
                _write_sample(out, outfull, f"{cik}:ALL", cik, g["trust_name"],
                              None, "", False, text, cand, ents)
                _fk = filter_grounded(cand, cik)
                if _fk:
                    n_fallback_trusts += 1
                    tgt = serialize_triples(_fk, ents)
                    n += 1; tot_text += len(text); tot_json += len(tgt)
                    ratios.append(round(len(text) / max(1, len(tgt)), 1))
                else:
                    n_dropped_empty += 1
                continue

            for label, seg in segs.items():
                f_iri = label_to_iri.get(label)
                if not f_iri:
                    continue
                cand = list(by_fund.get(f_iri, [])) + list(trust_edges)
                if not cand:
                    continue
                n_funds_located += 1
                cand, n_al, n_gnd, n_ctx = annotate_grounding(cand, ents, seg)
                sid = f"{cik}:{ents[f_iri].get('series_id') or _slug(label)}"
                _write_sample(out, outfull, sid, cik, g["trust_name"], label,
                              ents[f_iri].get("series_id", ""), True, seg, cand, ents)
                kept = filter_grounded(cand, cik)
                if kept:
                    tgt = serialize_triples(kept, ents)
                    n += 1; tot_text += len(seg); tot_json += len(tgt)
                    ratios.append(round(len(seg) / max(1, len(tgt)), 1))
                else:
                    n_dropped_empty += 1

    outfull.close()
    import statistics as _st
    log.info("wrote %d per-fund samples -> %s (full set -> %s)",
             n, SAMPLES_PATH, SAMPLES_FULL_PATH)
    if n:
        cov = n_funds_located / max(1, n_funds_total)
        print(f"\n{n} samples (grounding={GROUNDING_MODE}; "
              f"{n_funds_located}/{n_funds_total} funds located, coverage {cov:.0%}; "
              f"{n_fallback_trusts} whole-doc fallbacks; "
              f"{n_dropped_empty} units dropped with 0 grounded triples)")
        print(f"  mean input  : {tot_text//n:>8,} chars")
        print(f"  mean target : {tot_json//n:>8,} chars")
        print(f"  median ratio: {_st.median(ratios):>8.1f} : 1  (text : json)")
        print(f"  mean ratio  : {tot_text/max(1,tot_json):>8.1f} : 1")


# --------------------------------------------------------------------------- #
# stage: split  -- trust-level train/val/test split
# --------------------------------------------------------------------------- #
def _bucket(cik: str) -> float:
    """Deterministic [0,1) value per trust CIK (stable, no RNG state)."""
    import hashlib
    h = hashlib.sha256(cik.encode()).hexdigest()
    return int(h[:8], 16) / 0xFFFFFFFF


def build_split(val_frac: float = 0.10, test_frac: float = 0.10, src=None):
    """Split a sample file into train/val/test JSONL, partitioned by TRUST (CIK).

    Splitting by trust (not by fund) prevents leakage: two funds of the same
    trust share advisers, distributors and custodians, so allowing them into
    different splits would let the model memorise trust-specific entities. The
    assignment is a deterministic hash of the CIK, so the split is reproducible
    and stable as new samples are added. `src` defaults to samples.jsonl; pass
    trainset.jsonl to split the training-ready set.
    """
    src_path = Path(src) if src else SAMPLES_PATH
    if not src_path.exists():
        log.error("source not found: %s", src_path); return
    rows = [json.loads(l) for l in open(src_path, encoding="utf-8")]
    out = {"train": [], "val": [], "test": []}
    for r in rows:
        b = _bucket(r["cik"])
        split = "test" if b < test_frac else \
                "val" if b < test_frac + val_frac else "train"
        out[split].append(r)
    for name, recs in out.items():
        p = OUT_DIR / f"{name}.jsonl"
        with open(p, "w", encoding="utf-8") as f:
            for r in recs:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
    n_tr = {s: len({r["cik"] for r in recs}) for s, recs in out.items()}
    print("\nTrust-level split (deterministic by CIK):")
    for s in ("train", "val", "test"):
        print(f"  {s:5s}: {len(out[s]):>5,} samples  from {n_tr[s]:>4} trusts")
    total = sum(len(v) for v in out.values())
    print(f"  total: {total:,} samples -> {OUT_DIR}/{{train,val,test}}.jsonl")
    # leakage check
    cik_sets = {s: {r["cik"] for r in recs} for s, recs in out.items()}
    overlap = (cik_sets["train"] & cik_sets["val"]) | \
              (cik_sets["train"] & cik_sets["test"]) | \
              (cik_sets["val"] & cik_sets["test"])
    print(f"  trust overlap across splits: {len(overlap)} (should be 0)")


SAMPLES_LABELLED_PATH = OUT_DIR / "samples_labelled.jsonl"


def build_label():
    """Label extractability per triple and per sample, using the LLM match verdict.

    Reads the full annotated samples (SAMPLES_FULL_PATH, all gold triples with the
    cheap alias/grounded/context flags) and the LLM role-check verdicts
    (LLM_MATCH_PATH). For each triple it sets:
      - extractable        : the authoritative label = LLM role-check verdict
                             (the entity plays that role in the text);
      - alias_grounded/... : the cheap lexical flags (already present) for analysis.
    For each sample it sets:
      - n_extractable / n_triples
      - extractability = FULL  (all triples extractable)
                       | PARTIAL (some)
                       | NONE   (none)
    Writes:
      - SAMPLES_LABELLED_PATH (every sample, every triple labelled)
      - samples_full/partial/none.jsonl (samples split by label)
      - a distribution report (per label and per relation).
    """
    if not SAMPLES_FULL_PATH.exists():
        log.error("run `samples` first (need %s)", SAMPLES_FULL_PATH); return
    verdicts = _load_llm_verdicts()
    if not verdicts:
        log.error("no LLM match verdicts at %s — run `llm_extract.py --mode match` "
                  "first", LLM_MATCH_PATH); return

    by_label = {"FULL": [], "PARTIAL": [], "NONE": []}
    rel_tot = defaultdict(int); rel_ext = defaultdict(int)
    n = 0
    with open(SAMPLES_LABELLED_PATH, "w", encoding="utf-8") as out:
        for line in open(SAMPLES_FULL_PATH, encoding="utf-8"):
            r = json.loads(line)
            cik = r["cik"]
            triples = []
            n_ext = 0
            for t in r.get("target_triples", []):
                ext = verdicts.get((cik, t["p"], t["o"]), False)
                rel_tot[t["p"]] += 1
                rel_ext[t["p"]] += int(ext)
                n_ext += int(ext)
                triples.append({**t, "extractable": ext})
            ntot = len(triples)
            label = ("FULL" if ntot and n_ext == ntot else
                     "NONE" if n_ext == 0 else "PARTIAL")
            r["target_triples"] = triples
            r["extractability"] = label
            r.setdefault("stats", {})
            r["stats"]["n_extractable"] = n_ext
            r["stats"]["n_triples"] = ntot
            out.write(json.dumps(r, ensure_ascii=False) + "\n")
            by_label[label].append(r)
            n += 1

    for lab, recs in by_label.items():
        with open(OUT_DIR / f"labelled_{lab.lower()}.jsonl", "w", encoding="utf-8") as f:
            for r in recs:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

    tot_tr = sum(rel_tot.values()); ext_tr = sum(rel_ext.values())
    print(f"\nLabelled {n} samples -> {SAMPLES_LABELLED_PATH}")
    print(f"  sample extractability:")
    for lab in ("FULL", "PARTIAL", "NONE"):
        c = len(by_label[lab])
        print(f"    {lab:8s} {c:>5}  ({100*c/max(1,n):>3.0f}%)  -> labelled_{lab.lower()}.jsonl")
    print(f"  triple extractability: {ext_tr}/{tot_tr} = {100*ext_tr/max(1,tot_tr):.0f}%")
    print(f"\n  {'relation':16s}{'extractable':>12}{'%':>6}")
    for p in sorted(rel_tot, key=lambda x: -rel_tot[x]):
        print(f"  {p:16s}{str(rel_ext[p])+'/'+str(rel_tot[p]):>12}"
              f"{100*rel_ext[p]/rel_tot[p]:>5.0f}%")


TRAINSET_PATH = OUT_DIR / "trainset.jsonl"
EXCERPT_RADIUS = 600   # chars around each grounded name


_ROLE_HINT = ("custodian", "custody", "advis", "transfer agent", "administrat",
              "distributor", "underwriter", "serves as", "acts as", "series of")


def _evidence_excerpt(text_norm_full, raw_text, label, predicate=None):
    """Cut a coherent excerpt of raw_text around the BEST mention of the entity.

    A common firm token (e.g. "alliancebernstein") appears hundreds of times, so
    the first occurrence is usually boilerplate. We therefore score every
    occurrence of the name's most distinctive token by whether a role keyword
    (custodian/adviser/transfer agent/...) appears within +/- a window, and pick
    the occurrence with the nearest role hint -- that is the actual service-provider
    statement, the evidence we want. Falls back to the consecutive-tokens match,
    then the first token occurrence. Returns (excerpt, start) or (None, -1).
    """
    tk = _alias_tokens(label)
    if not tk:
        return None, -1
    low = raw_text.lower()
    distinctive = max(tk, key=len)
    occurrences = [m.start() for m in re.finditer(re.escape(distinctive), low)]
    pos = -1
    if occurrences:
        # prefer an occurrence with a role keyword nearby
        best = None
        for p in occurrences:
            win = low[max(0, p - 250): p + 250]
            if any(k in win for k in _ROLE_HINT):
                best = p
                break
        pos = best if best is not None else occurrences[0]
    if pos == -1:
        ac = _acronym(label)
        if len(ac) >= 3:
            m = re.search(r"\b" + re.escape(ac) + r"\b", low)
            pos = m.start() if m else -1
    if pos == -1:
        return None, -1
    a = max(0, pos - EXCERPT_RADIUS)
    b = min(len(raw_text), pos + EXCERPT_RADIUS)
    return raw_text[a:b].strip(), a


def build_trainset():
    """Assemble the training-ready dataset: per sample, the relevant TEXT EXCERPTS
    combined with ONLY the grounded (extractable) triples as the target.

    For each labelled sample (samples_labelled.jsonl):
      - keep only triples with extractable=True (every target is in the text);
      - for each, cut a coherent excerpt of the prospectus around the entity name
        and merge overlapping excerpts into a compact `input_text` (fits a model
        context window, unlike the multi-MB full book);
      - emit both serializations of the kept triples.
    Samples with no grounded triple (NONE) are skipped. Output: trainset.jsonl,
    then split into train/val/test by trust.
    """
    if not SAMPLES_LABELLED_PATH.exists():
        log.error("run `label` first (need %s)", SAMPLES_LABELLED_PATH); return
    gold = {json.loads(l)["cik"]: json.loads(l)
            for l in open(GOLD_PATH, encoding="utf-8")}
    n = n_skip = 0
    tot_in = tot_tr = 0
    with open(TRAINSET_PATH, "w", encoding="utf-8") as out:
        for line in open(SAMPLES_LABELLED_PATH, encoding="utf-8"):
            r = json.loads(line)
            ents = gold.get(r["cik"], {}).get("entities", {})
            full = r["input_text"]
            kept = [t for t in r["target_triples"] if t.get("extractable")]
            if not kept:
                n_skip += 1
                continue
            # collect excerpts around each grounded object name; merge overlaps
            spans = []
            for t in kept:
                lbl = ents.get(t["o"], {}).get("label", t["o"])
                exc, a = _evidence_excerpt(None, full, lbl)
                if exc:
                    spans.append((a, a + len(exc)))
            spans.sort()
            merged = []
            for a, b in spans:
                if merged and a <= merged[-1][1] + 200:
                    merged[-1] = (merged[-1][0], max(merged[-1][1], b))
                else:
                    merged.append((a, b))
            excerpt = "\n...\n".join(full[a:b].strip() for a, b in merged) or full[:2000]
            # restrict entities to kept triples
            ref = set()
            for t in kept:
                ref.add(t["s"]); ref.add(t["o"])
            se = {k: ents[k] for k in ref if k in ents}
            rec = {
                "sample_id": r["sample_id"], "cik": r["cik"],
                "trust_name": r.get("trust_name"),
                "input_text": excerpt,
                "ontology": ontology_schema(kept, se),
                "target_triples": kept,
                "target_serialized": serialize_triples(kept, se),
                "target_serialized_plain": serialize_triples_plain(kept, se),
                "stats": {"input_chars": len(excerpt), "n_triples": len(kept),
                          "text_to_json_ratio": round(
                              len(excerpt) / max(1, len(serialize_triples(kept, se))), 1)},
            }
            out.write(json.dumps(rec, ensure_ascii=False) + "\n")
            n += 1; tot_in += len(excerpt); tot_tr += len(kept)
    print(f"\nTrainset: {n} samples -> {TRAINSET_PATH}  ({n_skip} NONE-samples skipped)")
    if n:
        print(f"  grounded triples (targets): {tot_tr:,}")
        print(f"  mean input excerpt: {tot_in//n:,} chars  (was multi-MB full book)")
        print(f"  -> split with:  python build_rdf_dataset.py split --src trainset.jsonl")


# --------------------------------------------------------------------------- #
def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd")
    g = sub.add_parser("gold"); g.add_argument("--ncen", default="data/ncen/2025q3")
    g.add_argument("--custodian-scope", choices=["primary", "all", "none"],
                   default="primary",
                   help="primary=only IS_SUB_CUSTODIAN!=Y (default); all=every row; none=drop")
    f = sub.add_parser("fetch"); f.add_argument("--limit", type=int, default=25)
    f.add_argument("--max-filings", type=int, default=8,
                   help="max prospectus filings to fetch & concatenate per trust")
    f.add_argument("--no-sai", action="store_true", help="do not append SAI documents")
    f.add_argument("--ciks", nargs="*", help="only fetch these trust CIKs")
    s = sub.add_parser("samples")
    s.add_argument("--whole-trust", action="store_true",
                   help="one sample per trust instead of per-fund segmentation")
    s.add_argument("--grounding", choices=["alias", "llm", "name", "context", "none"], default="alias",
                   help="context=keep triples whose name+relation-keyword appear "
                        "(default, enforces text-extractability); name=name anywhere; "
                        "none=keep all, only annotate")
    sub.add_parser("label")     # label extractability (needs prior match run)
    sub.add_parser("trainset")  # assemble training set: excerpts + grounded triples
    sp = sub.add_parser("split")
    sp.add_argument("--val-frac", type=float, default=0.10)
    sp.add_argument("--test-frac", type=float, default=0.10)
    sp.add_argument("--src", default=None, help="source jsonl (default samples.jsonl; use trainset.jsonl)")
    a = sub.add_parser("all"); a.add_argument("--ncen", default="data/ncen/2025q3")
    a.add_argument("--limit", type=int, default=25)
    a.add_argument("--max-filings", type=int, default=8)
    a.add_argument("--custodian-scope", choices=["primary", "all", "none"], default="primary")
    a.add_argument("--whole-trust", action="store_true")
    a.add_argument("--grounding", choices=["alias", "llm", "name", "context", "none"], default="alias")
    a.add_argument("--val-frac", type=float, default=0.10)
    a.add_argument("--test-frac", type=float, default=0.10)
    args = ap.parse_args()

    global GROUNDING_MODE
    if args.cmd == "gold":
        build_gold(Path(args.ncen), custodian_scope=args.custodian_scope)
    elif args.cmd == "fetch":
        fetch_prose(args.limit, max_filings=args.max_filings,
                    include_sai=not args.no_sai, only_ciks=args.ciks)
    elif args.cmd == "samples":
        GROUNDING_MODE = args.grounding
        build_samples(per_fund=not args.whole_trust)
    elif args.cmd == "label":
        build_label()
    elif args.cmd == "trainset":
        build_trainset()
    elif args.cmd == "split":
        build_split(val_frac=args.val_frac, test_frac=args.test_frac, src=args.src)
    elif args.cmd == "all":
        GROUNDING_MODE = args.grounding
        build_gold(Path(args.ncen), custodian_scope=args.custodian_scope)
        fetch_prose(args.limit, max_filings=args.max_filings)
        build_samples(per_fund=not args.whole_trust)
        build_split(val_frac=args.val_frac, test_frac=args.test_frac)
    else:
        ap.print_help()


if __name__ == "__main__":
    main()