backend/scripts/extract_cuad_baselines.py

"""Compute empirical risk baselines from the CUAD corpus.

Runs the same feature extractors the risk engine judges against
(app/clause_features.py) over every CUAD contract, builds a distribution per
numeric feature, and derives the flag threshold from a percentile of that
distribution instead of a hand-typed constant. This is what makes the
baseline *data-derived* rather than asserted.

    # dry run — prints the computed distribution vs the current seed, writes nothing
    python -m scripts.extract_cuad_baselines            (from backend/)
    # promote: overwrite the numeric blocks of baselines_empirical.json
    python -m scripts.extract_cuad_baselines --write

Threshold policy (tunable): higher-worse features flag at the corpus p75
(top quartile is "above normal"); lower-worse features (uptime) flag at p25.
Requires CUAD at data/cuad/CUAD_v1.json (scripts/download_cuad.py).
"""
from __future__ import annotations

import json
import pathlib
import statistics
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from app.classification import classify_all  # noqa: E402
from app.clause_features import NUMERIC_EXTRACTORS  # noqa: E402
from app.schema import Clause, SourceSpan  # noqa: E402
from app.segmentation import CLAUSE_START  # noqa: E402

ROOT = pathlib.Path(__file__).resolve().parents[2]
CUAD_JSON = ROOT / "data" / "cuad" / "CUAD_v1.json"
TARGET = pathlib.Path(__file__).resolve().parents[1] / "app" / "baselines_empirical.json"

# higher-worse features flag at the top quartile; lower-worse at the bottom.
FLAG_PERCENTILE = {"higher_worse": 75, "lower_worse": 25}

# Careful-promotion guards: a computed threshold is only written if we have
# enough samples AND the value lands in a sane real-world range. This stops a
# noisy extractor or a tiny sample from silently weakening a real risk flag
# (e.g. breach-notice with n=3, or a polluted uptime distribution).
MIN_N = 15  # need a credible sample before a computed threshold overrides the vetted seed
SANE_RANGE = {
    ("liability_cap", "cap_months"): (1, 60),
    ("auto_renewal", "notice_days"): (7, 365),
    ("data_protection", "breach_notify_days"): (1, 30),
    ("payment_terms", "interest_pct_month"): (0.1, 10),
    ("payment_terms", "net_days"): (7, 180),
    ("sla", "uptime_pct"): (90, 100),
}


def _trustworthy(cat: str, feat: str, n: int, thr: float) -> tuple[bool, str]:
    if n < MIN_N:
        return False, f"n={n} < {MIN_N} (keep seed)"
    lo, hi = SANE_RANGE.get((cat, feat), (float("-inf"), float("inf")))
    if not (lo <= thr <= hi):
        return False, f"threshold {thr} outside sane [{lo},{hi}] (keep seed)"
    return True, "promote"


def segment_plain_text(text: str) -> list[Clause]:
    lines = text.split("\n")
    offsets, pos = [], 0
    for ln in lines:
        offsets.append(pos)
        pos += len(ln) + 1
    starts = [i for i, ln in enumerate(lines) if CLAUSE_START.match(ln)] or [0]
    out = []
    for n, li in enumerate(starts):
        s = offsets[li]
        e = offsets[starts[n + 1]] if n + 1 < len(starts) else len(text)
        out.append(Clause(id=f"c{n}", text=text[s:e],
                           span=SourceSpan(start_char=s, end_char=e, page=0)))
    return out


def pct(values: list[float], p: float) -> float:
    if not values:
        return 0.0
    s = sorted(values)
    k = (len(s) - 1) * (p / 100.0)
    lo, hi = int(k), min(int(k) + 1, len(s) - 1)
    return s[lo] + (s[hi] - s[lo]) * (k - lo)


def main() -> None:
    write = "--write" in sys.argv
    if not CUAD_JSON.exists():
        sys.exit(f"CUAD not found at {CUAD_JSON}. Run scripts/download_cuad.py first.")

    limit = 999
    for a in sys.argv:
        if a.startswith("--limit="):
            limit = int(a.split("=", 1)[1])

    data = json.loads(CUAD_JSON.read_text())["data"][:limit]
    samples: dict[tuple[str, str], list[float]] = {k: [] for k in NUMERIC_EXTRACTORS}

    for doc in data:
        for para in doc["paragraphs"]:
            clauses = segment_plain_text(para["context"])
            classify_all(clauses)
            for c in clauses:
                for (cat, feat), extract in NUMERIC_EXTRACTORS.items():
                    if cat in c.categories:
                        v = extract(c.text)
                        if v is not None:
                            samples[(cat, feat)].append(float(v))

    book = json.loads(TARGET.read_text())
    seed = {k: dict(v) for k, v in book.items()}  # for the diff print

    print(f"\nCUAD baseline computation — {len(data)} contracts\n")
    print(f"{'category.feature':<34} {'n':>5} {'median':>8} {'p75':>8} {'p90':>8}  "
          f"{'old→new threshold':>20}")
    print("-" * 92)
    for (cat, feat), vals in samples.items():
        node = book.setdefault(cat, {}).setdefault("numeric", {}).setdefault(feat, {})
        direction = node.get("direction", "higher_worse")
        old_thr = node.get("threshold")
        n = len(vals)
        if n == 0:
            print(f"{cat + '.' + feat:<34} {0:>5}  (no samples — keeping seed {old_thr})")
            continue
        dist = {
            "median": round(statistics.median(vals), 2),
            "p25": round(pct(vals, 25), 2), "p75": round(pct(vals, 75), 2),
            "p90": round(pct(vals, 90), 2),
            "min": round(min(vals), 2), "max": round(max(vals), 2),
        }
        new_thr = round(pct(vals, FLAG_PERCENTILE[direction]), 2)
        ok, why = _trustworthy(cat, feat, n, new_thr)
        flag = "✓ promote" if ok else f"✗ {why}"
        print(f"{cat + '.' + feat:<34} {n:>5} {dist['median']:>8} {dist['p75']:>8} "
              f"{dist['p90']:>8}  {str(old_thr):>9} → {new_thr:<7} {flag}")
        if write and ok:
            # promote threshold + distribution (drives provenance and severity
            # scaling). Untrusted features are left pure-seed (no n/dist), so a
            # thin sample or noisy extractor can never weaken a real flag.
            node.update({"threshold": new_thr, "n": n, "dist": dist})

    if write:
        book.setdefault("_meta", {}).update({"computed": True})
        TARGET.write_text(json.dumps(book, indent=2) + "\n")
        print(f"\nwrote computed baselines to {TARGET} (computed=true)")
    else:
        print("\ndry run — pass --write to overwrite baselines_empirical.json")
        _ = seed  # noqa


if __name__ == "__main__":
    main()