#!/usr/bin/env python3
"""Independent second read of F2. Re-implements corpus loading, URL de-duplication, near-duplicate clustering,
per-year counting and dollar sums without importing reproduce_enforcement.py or load_corpus.py. The only shared
code is the frozen rules module (handcheck/doj_rules_FROZEN.py), which is the classification definition that the
hand-check validated. Compares every cell of final/doj-pandemic-enforcement-by-year.csv.
Usage: check_enforcement.py <raw_doj_dir> <topup.jsonl.gz> <by-year.csv>"""
import csv, datetime, gzip, html, importlib.util, json, os, re, sys
from collections import defaultdict

HERE = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location("rules", os.path.join(HERE, "handcheck", "doj_rules_FROZEN.py"))
rules = importlib.util.module_from_spec(spec); spec.loader.exec_module(rules)

raw, topup, table = sys.argv[1], sys.argv[2], sys.argv[3]
paths = [os.path.join(raw, "doj_press_releases_2020-2026_slim.jsonl.gz"),
         os.path.join(raw, "doj_press_releases_2026-07-04_2026-09-15_slim.jsonl.gz"), topup]
latest = {}
for p in paths:                      # later files override earlier copies of the same URL
    with gzip.open(p, "rt", encoding="utf-8") as fh:
        for ln in fh:
            rec = json.loads(ln)
            key = (rec.get("url") or "").strip().rstrip("/").lower()
            if key:
                latest[key] = rec
year_all = defaultdict(int); kept = []
for rec in latest.values():
    try:
        d = datetime.datetime.utcfromtimestamp(int(rec.get("date"))).date()
    except Exception:
        continue
    year_all[str(d.year)] += 1
    t, b = (rec.get("title") or "").strip(), rec.get("body_text") or ""
    if rules.is_translation(t) or not rules.scope(t, b):
        continue
    st = rules.stage(t, b)[0]
    crim = st in ("CHARGED", "PLEA", "CONVICTED", "SENTENCED")
    amt = rules.settlement_usd(t, b)[0] if st == "CIVIL_RESOLUTION" else None
    toks = html.unescape(t).lower()
    toks = re.sub(r"\$\s?([\d.,]+)\s*m\b", r"$\1 million", toks)
    toks = re.sub(r"\$\s?([\d.,]+)\s*b\b", r"$\1 billion", toks).replace("payment protection", "paycheck protection").replace("&#039;", "'")
    toks = set(re.sub(r"\s+", " ", re.sub(r"[^a-z0-9$. ]", " ", toks)).split())
    kept.append({"d": d, "u": rec.get("url"), "y": str(d.year), "st": st, "amt": amt,
                 "def": rules.defendants_estimate(t, b) if crim else None, "tok": toks,
                 "lead": re.sub(r"\s+", " ", b)[:300]})
kept.sort(key=lambda k: (k["d"].isoformat(), k["u"]))
# greedy clustering, same stated rule: same stage; within 30 days of a cluster head; title-token Jaccard >= 0.8,
# or identical first 300 body characters, or (civil) same amount within 0.1% inside 3 days or within 2% inside 7 days
# with Jaccard >= 0.25. A release joins the most recent qualifying head, else starts a cluster.
heads = defaultdict(list)
for k in kept:
    k["dup"] = False
    for h in reversed(heads[k["st"]]):
        gap = (k["d"] - h["d"]).days
        if gap > 30:
            break
        jac = len(k["tok"] & h["tok"]) / max(1, len(k["tok"] | h["tok"]))
        same = False
        if k["st"] == "CIVIL_RESOLUTION" and k["amt"] is not None and h["amt"] is not None:
            big = max(k["amt"], h["amt"]); diff = abs(k["amt"] - h["amt"])
            same = (gap <= 3 and diff <= 0.001 * big) or (gap <= 7 and diff <= 0.02 * big and jac >= 0.25)
        if jac >= 0.8 or (k["lead"] and k["lead"] == h["lead"]) or same:
            k["dup"] = True; break
    if not k["dup"]:
        heads[k["st"]].append(k)
agg = defaultdict(lambda: {"releases": 0, "distinct": 0, "def": 0, "defcap": 0, "usd": 0.0, "n_amt": 0})
for k in kept:
    a = agg[(k["y"], k["st"])]; a["releases"] += 1
    if k["dup"]:
        continue
    a["distinct"] += 1
    if k["def"] is not None:
        a["def"] += k["def"]; a["defcap"] += min(k["def"], 20)
    if k["amt"] is not None:
        a["usd"] += k["amt"]; a["n_amt"] += 1
ok = bad = 0
for r in csv.DictReader(open(table, encoding="utf-8")):
    a = agg[(r["year"], r["stage"])]
    pairs = [("releases", a["releases"], int(r["releases"])), ("distinct", a["distinct"], int(r["distinct_announcements"])),
             ("all_doj", year_all[r["year"]], int(r["all_doj_releases_that_year"]))]
    if r["defendants_estimate"] != "":
        pairs += [("defendants", a["def"], int(r["defendants_estimate"])), ("defendants_cap20", a["defcap"], int(r["defendants_estimate_capped20"]))]
    if r["stage"] == "CIVIL_RESOLUTION":
        pairs += [("usd", round(a["usd"], 2), round(float(r["settlement_usd_sum"]), 2)), ("n_amt", a["n_amt"], int(r["settlements_with_amount"]))]
    for name, x, y in pairs:
        if abs(x - y) <= 0.01:
            ok += 1
        else:
            bad += 1; print(f"FAIL {r['year']} {r['stage']} {name}: check={x} reproduce={y}")
print(f"RESULT {ok} PASS / {bad} FAIL (cells compared across {len(set(k for k, _ in agg))} years)")