"""N5 - focus-window coverage: what share of every ledger section the model actually saw (cluster E).

Imports _focus / _PROJECT_GATE from the hash-verified snapshot (pure functions). For each of the
1,123,098 ledger sections: length, Python gate pass, gate hits, shown characters (replicated span
logic, asserted equal to the real _focus output), offset of the last shown character, hits shown,
14,000-character message cut, capex vocabulary inside/outside the shown windows. Aggregates by form x
item; n_mentions vs length (plateau); <30-word / incorporated-by-reference 10-K MD&As by year;
median words of ledger 10-K MD&A (the 9,905 figure). Read-only; per-section features kept on the node.
"""
from __future__ import annotations
import json, re, sys, time
from pathlib import Path
import multiprocessing as mp
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parent))
from spid.extract.llm import _focus          # snapshot copy
from spid.extract.rules import _PROJECT_GATE  # snapshot copy

BUDGET = 9000
CAPEX = re.compile(r"capital expenditure|liquidity and capital resources|commitment", re.I)

def spans_like_focus(text: str, hits):
    """Replicates _focus (llm.py:55-80) and returns the chunks [(a,b)] actually emitted."""
    L = len(text)
    if L <= BUDGET:
        return [(0, L)], True
    spans = []
    for s, e in hits:
        spans.append((max(0, s - 220), min(L, e + 380)))
        if sum(b - a for a, b in spans) > BUDGET * 1.4:
            break
    if not spans:
        return [(0, BUDGET)], True
    spans.sort()
    merged = [spans[0]]
    for a, b in spans[1:]:
        if a <= merged[-1][1]:
            merged[-1] = (merged[-1][0], max(merged[-1][1], b))
        else:
            merged.append((a, b))
    chunks, total = [], 0
    for a, b in merged:
        chunks.append((a, b)); total += b - a
        if total >= BUDGET:
            break
    return chunks, False

def feat(item):
    sid, text = item
    L = len(text)
    hits = [(m.start(), m.end()) for m in _PROJECT_GATE.finditer(text)]
    nh = len(hits)
    chunks, whole = spans_like_focus(text, hits)
    shown = sum(b - a for a, b in chunks)
    last = chunks[-1][1]
    focus = _focus(text)
    eq = ("\n...\n".join(text[a:b] for a, b in chunks) == focus)
    # hits inside the chunks
    hs = 0; ci = 0
    for s, _ in hits:
        while ci < len(chunks) and s >= chunks[ci][1]:
            ci += 1
        if ci < len(chunks) and chunks[ci][0] <= s < chunks[ci][1]:
            hs += 1
    cap_all = 0; cap_in = 0
    if L > BUDGET:
        for m in CAPEX.finditer(text):
            cap_all += 1
            s = m.start()
            if any(a <= s < b for a, b in chunks):
                cap_in += 1
    return (sid, L, nh > 0, nh, shown, last, hs, len(focus), eq, cap_all, cap_in, len(chunks))

def main():
    r = Run("focus")
    con = connect(SPID_DB)
    con.execute("CREATE TEMP TABLE led AS SELECT section_id FROM extracted_sections")
    q = """SELECT s.section_id, s.text FROM src_sections s JOIN led USING (section_id)
           UNION ALL SELECT s.section_id, s.text FROM spid_sections s JOIN led USING (section_id)"""
    cur = con.execute(q)
    rows_out = []; t0 = time.time(); done = 0; chars = 0
    ctx = mp.get_context("fork")
    with ctx.Pool(4) as pool:
        while True:
            rows = cur.fetchmany(4000)
            if not rows:
                break
            chars += sum(len(t) for _, t in rows)
            rows_out.extend(pool.imap_unordered(feat, rows, chunksize=100))
            done += len(rows)
            if done % 100000 < 4000:
                print(f"  {done} sections, {chars/1e9:.2f} GB, {time.time()-t0:.0f}s", flush=True)
    r.notes["sections_streamed"] = done; r.notes["chars_streamed"] = chars
    import pyarrow as pa, pyarrow.parquet as pq
    names = ["section_id", "nchar", "gate_py", "n_hits", "shown_chars", "last_offset", "hits_shown", "focus_len", "replica_eq", "capex_all", "capex_shown", "n_chunks"]
    tbl = pa.table({n: [row[i] for row in rows_out] for i, n in enumerate(names)})
    pq.write_table(tbl, r.out / "section_features.parquet")   # stays on the node
    con.register("ft", tbl)
    tmin = con.execute("select min(extracted_at) from extracted_sections").fetchone()[0]
    con.execute(f"""CREATE TEMP TABLE F AS
    SELECT ft.*, e.n_mentions, CASE WHEN e.extracted_at = TIMESTAMP '{tmin}' THEN 'backfilled_june' ELSE 'october' END run,
           m.section_name item, m.word_count, f.form_type form, year(f.filing_date) yr,
           shown_chars*1.0/nchar AS shown_frac, last_offset*1.0/nchar AS last_frac,
           CASE WHEN n_hits > 0 THEN hits_shown*1.0/n_hits END AS hits_frac,
           (9 + length(m.section_name) + 2 + focus_len) > 14000 AS msg_truncated,
           nchar > {BUDGET} AS over_budget
    FROM ft JOIN extracted_sections e USING (section_id)
    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)
    JOIN src_filings f ON f.accession_number = e.accession_number""")
    r.notes["replica_mismatches"] = con.execute("select count(*) from F where not replica_eq").fetchone()[0]
    r.notes["gate_py_pass"] = con.execute("select count(*) from F where gate_py").fetchone()[0]
    r.notes["gate_py_pass_by_form"] = con.execute("select form, count(*) filter (where gate_py), count(*) from F group by 1 order by 1").fetchall()
    r.notes["gate_fail_with_mentions"] = con.execute("select count(*) from F where not gate_py and n_mentions>0").fetchone()[0]
    agg = """count(*) n, count(*) FILTER (WHERE gate_py) gate_pass, count(*) FILTER (WHERE over_budget) over_9000_chars, round(100.0*count(*) FILTER (WHERE over_budget)/count(*),2) pct_over_9000,
      round(quantile_cont(nchar,0.1)) chars_p10, round(quantile_cont(nchar,0.5)) chars_p50, round(quantile_cont(nchar,0.9)) chars_p90,
      round(quantile_cont(shown_frac,0.1),3) shown_p10, round(quantile_cont(shown_frac,0.5),3) shown_p50, round(quantile_cont(shown_frac,0.9),3) shown_p90, round(avg(shown_frac),3) shown_mean,
      round(quantile_cont(shown_frac,0.5) FILTER (WHERE over_budget),3) shown_p50_over_budget,
      round(quantile_cont(last_frac,0.1),3) last_p10, round(quantile_cont(last_frac,0.5),3) last_p50, round(quantile_cont(last_frac,0.9),3) last_p90,
      round(100.0*count(*) FILTER (WHERE shown_frac < 0.25)/count(*),2) pct_shown_lt_25, round(100.0*count(*) FILTER (WHERE last_frac < 0.5)/count(*),2) pct_last_lt_50,
      round(100.0*count(*) FILTER (WHERE msg_truncated)/count(*),2) pct_msg_truncated_14k, round(avg(hits_frac),3) hits_shown_mean, round(quantile_cont(hits_frac,0.5),3) hits_shown_p50,
      sum(n_mentions) mentions, sum(n_mentions) FILTER (WHERE shown_frac < 0.25) mentions_from_lt25, round(100.0*sum(n_mentions) FILTER (WHERE shown_frac < 0.25)/nullif(sum(n_mentions),0),2) pct_mentions_from_lt25,
      round(100.0*sum(capex_shown)/nullif(sum(capex_all),0),2) pct_capex_vocab_shown_over_budget"""
    r.dump("focus_coverage_by_form_item.csv", con, f"SELECT form, item, {agg} FROM F GROUP BY 1,2 ORDER BY 1,2")
    r.dump("focus_coverage_by_form.csv", con, f"SELECT form, {agg} FROM F GROUP BY 1 ORDER BY 1")
    r.dump("focus_coverage_by_form_item_run.csv", con, f"SELECT form, item, run, {agg} FROM F GROUP BY 1,2,3 ORDER BY 1,2,3")
    r.dump("focus_coverage_calls_only_by_form_item.csv", con, f"SELECT form, item, {agg} FROM F WHERE gate_py AND word_count >= 30 GROUP BY 1,2 ORDER BY 1,2")
    r.dump("mentions_vs_length_by_item.csv", con, """
    SELECT form, item, CASE WHEN nchar<3000 THEN 'a<3k' WHEN nchar<6000 THEN 'b3-6k' WHEN nchar<9000 THEN 'c6-9k' WHEN nchar<12000 THEN 'd9-12k' WHEN nchar<20000 THEN 'e12-20k' WHEN nchar<40000 THEN 'f20-40k' WHEN nchar<80000 THEN 'g40-80k' ELSE 'h>=80k' END len_bin,
           count(*) n, count(*) FILTER (WHERE gate_py) calls, round(avg(n_mentions),4) mean_mentions, round(avg(n_mentions) FILTER (WHERE gate_py),4) mean_mentions_per_call,
           round(100.0*count(*) FILTER (WHERE n_mentions>0)/count(*),2) productive_pct, round(avg(shown_frac),3) shown_mean, round(avg(shown_chars)) shown_chars_mean,
           round(1e4*sum(n_mentions)/nullif(sum(shown_chars),0),4) mentions_per_10k_shown_chars
    FROM F GROUP BY 1,2,3 ORDER BY 1,2,3""")
    r.dump("shown_fraction_bins.csv", con, """
    SELECT form, CASE WHEN shown_frac >= 0.9 THEN 'a>=0.9' WHEN shown_frac >= 0.25 THEN 'b0.25-0.9' ELSE 'c<0.25' END bin, count(*) n, count(*) FILTER (WHERE gate_py AND word_count>=30) calls,
           count(*) FILTER (WHERE n_mentions>0) with_mentions, sum(n_mentions) mentions FROM F GROUP BY 1,2 ORDER BY 1,2""")
    # 10-K MD&A coverage hole by year: all parsed + legacy 10-K md_and_a sections (ledgered or not)
    r.dump("mdna_10k_short_by_year.csv", con, """
    WITH a AS (SELECT s.section_id, s.word_count, year(f.filing_date) yr, s.section_id IN (SELECT section_id FROM extracted_sections) ledgered,
                      regexp_matches(left(s.text, 3000), 'incorporated .{0,40}by reference', 'i') incorp_head
               FROM src_sections s JOIN src_filings f USING (accession_number) WHERE f.form_type='10-K' AND s.section_name='md_and_a' AND s.accession_number NOT IN (SELECT accession_number FROM spid_sections)
               UNION ALL
               SELECT s.section_id, s.word_count, year(f.filing_date), TRUE, regexp_matches(left(s.text, 3000), 'incorporated .{0,40}by reference', 'i')
               FROM spid_sections s JOIN src_filings f USING (accession_number) WHERE s.section_name='md_and_a')
    SELECT yr, count(*) mdna_sections, count(*) FILTER (WHERE word_count < 30) lt30_words, round(100.0*count(*) FILTER (WHERE word_count < 30)/count(*),2) pct_lt30,
           count(*) FILTER (WHERE word_count < 300 AND incorp_head) lt300_incorporated_by_ref, count(*) FILTER (WHERE NOT ledgered) not_ledgered, round(median(word_count)) median_words
    FROM a GROUP BY 1 ORDER BY 1""")
    K = {}
    K["ledger_10k_mdna_median_words"] = con.execute("select median(word_count) from F where form='10-K' and item='md_and_a'").fetchone()[0]
    K["ledger_10k_mdna_median_chars"] = con.execute("select median(nchar) from F where form='10-K' and item='md_and_a'").fetchone()[0]
    K["ledger_10k_mdna_n"] = con.execute("select count(*) from F where form='10-K' and item='md_and_a'").fetchone()[0]
    K["ledger_10k_mdna_shown_p50"] = con.execute("select median(shown_frac) from F where form='10-K' and item='md_and_a'").fetchone()[0]
    K["ledger_10k_mdna_last_p50"] = con.execute("select median(last_frac) from F where form='10-K' and item='md_and_a'").fetchone()[0]
    K["ledger_10k_mdna_pct_shown_lt_25"] = con.execute("select 100.0*count(*) filter (where shown_frac<0.25)/count(*) from F where form='10-K' and item='md_and_a'").fetchone()[0]
    K["ledger_sections_over_9000_chars"], K["ledger_sections"] = con.execute("select count(*) filter (where over_budget), count(*) from F").fetchone()
    K["calls_over_9000_chars"], K["calls"] = con.execute("select count(*) filter (where over_budget), count(*) from F where gate_py and word_count>=30").fetchone()
    K["mentions_from_sections_shown_lt_25"], K["mentions_total"] = con.execute("select sum(n_mentions) filter (where shown_frac<0.25), sum(n_mentions) from F").fetchone()
    K["msg_truncated_14k"] = con.execute("select count(*) from F where msg_truncated").fetchone()[0]
    K["gate_py_pass_total"] = r.notes["gate_py_pass"]; K["gate_duckdb_pass_total_N1"] = 703737
    K["mean_shown_frac_over_budget"] = con.execute("select avg(shown_frac) from F where over_budget").fetchone()[0]
    r.write_json("stats_keys.json", K)
    r.finish()

if __name__ == "__main__":
    main()