"""N8 - rule-based baseline (snapshot rules.py) vs the LLM output on a stratified seeded sample of ledger
sections, plus rules-on-full vs rules-on-focused text (recall-loss bound, THEO-09). Read-only; CPU only.
"""
from __future__ import annotations
import json, sys, time
from pathlib import Path
import multiprocessing as mp
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import *
from spid.extract.rules import RuleExtractor
from spid.extract.llm import _focus

N_PROD = 30000; N_ZERO = 30000
RX = RuleExtractor()

def work(item):
    sid, name, text = item
    full = RX.extract(text=text, accession_number="a", cik="c", section_id=sid, section_name=name)
    foc = RX.extract(text=_focus(text), accession_number="a", cik="c", section_id=sid, section_name=name)
    return (sid, len(full), sorted({m.project_type for m in full}), len(foc), sorted({m.project_type for m in foc}), len(text),
            [(m.project_type, m.confidence, m.amount_usd, m.description[:200]) for m in full][:20])

def main():
    r = Run("rules")
    con = connect(SPID_DB)
    tmin = con.execute("select min(extracted_at) from extracted_sections").fetchone()[0]
    con.execute(f"""CREATE TEMP TABLE led AS SELECT e.section_id, e.n_mentions, CASE WHEN e.extracted_at = TIMESTAMP '{tmin}' THEN 'backfilled_june' ELSE 'october' END run,
       f.form_type, m.section_name, m.word_count FROM extracted_sections e JOIN src_filings f ON f.accession_number = e.accession_number
       JOIN (SELECT section_id, section_name, word_count FROM src_sections UNION ALL SELECT section_id, section_name, word_count FROM spid_sections) m USING (section_id)""")
    N_prod_pop = con.execute("select count(*) from led where n_mentions>0").fetchone()[0]
    N_zero_pop = con.execute("select count(*) from led where n_mentions=0").fetchone()[0]
    # DuckDB applies USING SAMPLE before WHERE: the stratum filter must sit in a subquery
    con.execute(f"CREATE TEMP TABLE samp AS SELECT *, 'productive' stratum, {N_prod_pop}/{N_PROD}.0 w FROM (SELECT * FROM led WHERE n_mentions>0) USING SAMPLE reservoir({N_PROD} ROWS) REPEATABLE ({SEED})")
    con.execute(f"INSERT INTO samp SELECT *, 'zero', {N_zero_pop}/{N_ZERO}.0 FROM (SELECT * FROM led WHERE n_mentions=0) USING SAMPLE reservoir({N_ZERO} ROWS) REPEATABLE ({SEED})")
    assert con.execute("select count(*) filter (where n_mentions>0), count(*) filter (where n_mentions=0) from samp").fetchone() == (N_PROD, N_ZERO)
    n_s = con.execute("select count(*) from samp").fetchone()[0]
    print(f"sample {n_s} sections (pop productive {N_prod_pop}, zero {N_zero_pop})", flush=True)
    cur = con.execute("""SELECT s.section_id, s.section_name, t.text FROM samp s JOIN (SELECT section_id, text FROM src_sections UNION ALL SELECT section_id, text FROM spid_sections) t USING (section_id)""")
    out = []; t0 = time.time(); done = 0
    ctx = mp.get_context("fork")
    with ctx.Pool(4) as pool:
        while True:
            rows = cur.fetchmany(2000)
            if not rows: break
            out.extend(pool.imap_unordered(work, rows, chunksize=25))
            done += len(rows)
            print(f"  {done}/{n_s} sections, {time.time()-t0:.0f}s", flush=True)
    import pyarrow as pa, pyarrow.parquet as pq
    tbl = pa.table({"section_id": [o[0] for o in out], "rules_n": [o[1] for o in out], "rules_types": [o[2] for o in out], "rules_focus_n": [o[3] for o in out],
                    "rules_focus_types": [o[4] for o in out], "nchar": [o[5] for o in out], "rules_mentions": [json.dumps(o[6]) for o in out]})
    pq.write_table(tbl, r.out / "rules_mentions.parquet")   # stays on the node
    con.register("rt", tbl)
    con.execute("""CREATE TEMP TABLE lt AS SELECT section_id, list_distinct(list(project_type)) llm_types FROM project_mentions GROUP BY 1""")
    con.execute("""CREATE TEMP TABLE J AS SELECT s.*, r.rules_n, r.rules_types, r.rules_focus_n, r.rules_focus_types, r.nchar, coalesce(l.llm_types, []) llm_types,
        r.rules_n > 0 rules_pos, s.n_mentions > 0 llm_pos, len(list_intersect(r.rules_types, coalesce(l.llm_types, []))) > 0 type_overlap
        FROM samp s JOIN rt r USING (section_id) LEFT JOIN lt l USING (section_id)""")
    agg = """count(*) n_sampled, round(sum(w)) n_weighted,
      count(*) FILTER (WHERE llm_pos) llm_pos, count(*) FILTER (WHERE rules_pos) rules_pos,
      round(100.0*count(*) FILTER (WHERE llm_pos AND rules_pos)/nullif(count(*) FILTER (WHERE llm_pos),0),2) p_rules_given_llm,
      round(100.0*sum(w) FILTER (WHERE llm_pos AND rules_pos)/nullif(sum(w) FILTER (WHERE rules_pos),0),2) p_llm_given_rules_weighted,
      round(100.0*sum(w) FILTER (WHERE rules_pos)/sum(w),2) rules_pos_pct_weighted, round(100.0*sum(w) FILTER (WHERE llm_pos)/sum(w),2) llm_pos_pct_weighted,
      round(100.0*count(*) FILTER (WHERE llm_pos AND rules_pos AND type_overlap)/nullif(count(*) FILTER (WHERE llm_pos AND rules_pos),0),2) type_overlap_pct_codetected,
      round(100.0*sum(w) FILTER (WHERE llm_pos AND NOT rules_pos)/sum(w),2) llm_only_pct_w, round(100.0*sum(w) FILTER (WHERE rules_pos AND NOT llm_pos)/sum(w),2) rules_only_pct_w,
      round(100.0*sum(w) FILTER (WHERE llm_pos = rules_pos)/sum(w),2) agreement_pct_w,
      round(avg(rules_n) FILTER (WHERE llm_pos),3) rules_mentions_per_llm_pos_section, round(avg(rules_n) FILTER (WHERE NOT llm_pos),3) rules_mentions_per_llm_zero_section"""
    r.dump("agreement_by_form_item.csv", con, f"SELECT form_type, section_name item, {agg} FROM J GROUP BY 1,2 ORDER BY 1,2")
    r.dump("agreement_by_form.csv", con, f"SELECT form_type, {agg} FROM J GROUP BY 1 ORDER BY 1")
    r.dump("agreement_by_run.csv", con, f"SELECT run, {agg} FROM J GROUP BY 1 ORDER BY 1")
    r.dump("agreement_overall.csv", con, f"SELECT 'all' scope, {agg} FROM J")
    r.dump("type_agreement.csv", con, """
      WITH x AS (SELECT unnest(llm_types) t, rules_types FROM J WHERE llm_pos AND rules_pos)
      SELECT t llm_type, count(*) n, count(*) FILTER (WHERE list_contains(rules_types, t)) same_type_in_rules, round(100.0*count(*) FILTER (WHERE list_contains(rules_types, t))/count(*),2) pct FROM x GROUP BY 1 ORDER BY 2 DESC""")
    r.dump("focus_recall_loss.csv", con, """
      SELECT form_type, section_name item, count(*) FILTER (WHERE nchar > 9000) n_over_budget,
             sum(rules_n) FILTER (WHERE nchar > 9000) rules_full, sum(rules_focus_n) FILTER (WHERE nchar > 9000) rules_focused,
             round(100.0*(1 - sum(rules_focus_n) FILTER (WHERE nchar > 9000)*1.0/nullif(sum(rules_n) FILTER (WHERE nchar > 9000),0)),2) pct_rules_mentions_lost_to_focus,
             round(100.0*count(*) FILTER (WHERE nchar > 9000 AND rules_n > 0 AND rules_focus_n = 0)/nullif(count(*) FILTER (WHERE nchar > 9000 AND rules_n > 0),0),2) pct_sections_lost_entirely,
             round(100.0*sum(w) FILTER (WHERE nchar > 9000 AND rules_n > 0 AND rules_focus_n = 0)/nullif(sum(w) FILTER (WHERE nchar > 9000 AND rules_n > 0),0),2) pct_sections_lost_entirely_w
      FROM J GROUP BY 1,2 ORDER BY 1,2""")
    r.dump("focus_recall_loss_overall.csv", con, """
      SELECT 'over_9000_chars' scope, count(*) n, sum(rules_n) rules_full, sum(rules_focus_n) rules_focused, round(100.0*(1 - sum(rules_focus_n)*1.0/nullif(sum(rules_n),0)),2) pct_lost,
             round(100.0*sum(w*rules_n)/nullif(sum(w*rules_n) + 0,0),2) dummy FROM J WHERE nchar > 9000""")
    r.dump("disagreement_strata.csv", con, """SELECT form_type, section_name item, run, llm_pos, rules_pos, count(*) n_sampled, round(sum(w)) n_weighted FROM J GROUP BY 1,2,3,4,5 ORDER BY 1,2,3,4,5""")
    r.notes["sample"] = {"productive": N_PROD, "zero": N_ZERO, "pop_productive": N_prod_pop, "pop_zero": N_zero_pop, "seed": SEED}
    r.finish()

if __name__ == "__main__":
    main()