#!/usr/bin/env python3
"""exp16 — cold-start analysis: spread + canon-recovery (confound removed) + A4/A5 banknote (research/20
§D1 cold-start, §D2 A4/A5, §D5 decision rules). $0 (analysis of the paid cold-start drafts).

Compares the injection-confounded records.json drafts (§D1 lower bound) against the cold-start drafts
(no glossary injection — ecological W1.5 input) on the SAME chapters [1,4,5,7,9]:
  • spread on GT terms: injected ≈0 (suppressed) vs cold-start (the real dispersion signal)
  • canon-recovery: does §C2 recover the owner-signed dst from an UN-anchored draft? (§D5 hypothesis-1 gate)
  • A4/A5 banknote: marginal recall on f<3, parse_fail/truncated rates, quality delta (plain vs banknote),
    real token cost. §D5 rule applied verbatim.
"""
from __future__ import annotations

import json
from pathlib import Path

import exp16_common as X
import spread as SP
import canon as CANON
import detectors as D
import patterns as P
import arms as A

BOOK = Path("/home/ubuntu/books/gu-zhenren")
COLD = BOOK / "exp16" / "coldstart"
COLD_CHAPTERS = [1, 4, 5, 7, 9]


def load_coldstart_chunks(config="plain"):
    """Build Chunk objects: source from records.json, draft from the cold-start config output."""
    recs = {(r["chapter"], r["chunk_idx"]): r for r in json.load(open(BOOK / "rerun" / "records.json"))}
    out = []
    src_dir = COLD / config
    for (ch, ck), r in recs.items():
        if ch not in COLD_CHAPTERS:
            continue
        raw = src_dir / (f"{ch}.{ck}.clean.txt" if config == "banknote" else f"{ch}.{ck}.raw.txt")
        if not raw.exists():
            continue
        c = X.Chunk(chapter=ch, chunk_idx=ck, source=r["source"], draft=raw.read_text(encoding="utf-8"))
        c.nsource = X.norm(c.source)
        out.append(c)
    out.sort(key=lambda c: (c.chapter, c.chunk_idx))
    return out


def injected_chunks_for_cold():
    all_ch = X.load_chunks()
    return [c for c in all_ch if c.chapter in COLD_CHAPTERS]


def gt_in(chunks, gt):
    return [e for e in gt if X.gt_occurrences(e, chunks) >= 1]


def spread_compare(gt, inj_chunks, cold_chunks):
    sm_inj = SP.SpreadModel(inj_chunks)
    sm_cold = SP.SpreadModel(cold_chunks)
    rows = []
    for e in gt:
        occ = X.gt_occurrences(e, cold_chunks)
        if occ < 2:      # spread needs >=2 occurrences
            continue
        sn = X.norm(e.src)
        rows.append(dict(src=e.src, typ=e.typ, occ=occ,
                         spread_injected=round(sm_inj.spread(sn), 3),
                         spread_coldstart=round(sm_cold.spread(sn), 3)))
    return rows, sm_inj, sm_cold


def canon_compare(gt, inj_chunks, cold_chunks):
    gt_cold = gt_in(cold_chunks, gt)
    rows_inj, rate_inj, tot_inj = CANON.run(inj_chunks, gt_cold)
    rows_cold, rate_cold, tot_cold = CANON.run(cold_chunks, gt_cold)
    return dict(injected=(rate_inj, tot_inj), coldstart=(rate_cold, tot_cold)), rows_cold


def banknote_analysis(gt):
    """A4/A5: aggregate banknote entries, marginal recall on f<3, parse_fail/truncated, quality delta."""
    bdir = COLD / "banknote"
    entries, flags_all = [], []
    per_chunk = []
    for f in sorted(bdir.glob("*.banknote.json")):
        d = json.load(open(f))
        entries += d["entries"]
        flags_all.append(d["flags"])
        per_chunk.append((f.stem, d["flags"]))
    proposed = {X.norm(e["src"]) for e in entries}
    # recall of GT by banknote channel, by stratum
    all_chunks = X.load_chunks()
    def bank_recall(stratum=None):
        hit = tot = 0
        for e in gt:
            f = X.gt_occurrences(e, all_chunks)
            if stratum and X.freq_stratum(f) != stratum:
                continue
            tot += 1
            hit += any(s in proposed for s in e.norm_surfaces)
        return hit, tot
    n_lines = sum(fl["n_banknote_lines"] for fl in flags_all)
    n_parse_fail = sum(fl["banknote_parse_fail"] for fl in flags_all)
    n_trunc = sum(fl["banknote_truncated"] for fl in flags_all)
    nchunks = len(flags_all)
    return dict(n_entries=len(entries), distinct_src=len(proposed),
                recall_all=bank_recall(), recall_f_lt3=bank_recall("f<3"),
                recall_f3_9=bank_recall("f3-9"), recall_f_ge10=bank_recall("f>=10"),
                n_banknote_lines=n_lines, parse_fail_chunks=n_parse_fail, truncated_chunks=n_trunc,
                nchunks=nchunks, parse_fail_rate=(n_parse_fail / nchunks if nchunks else 0),
                sample_entries=entries[:20]), proposed


def quality_delta():
    """Deterministic draft-quality delta plain vs banknote (length ratio, char-count) — does the banknote
    instruction degrade the translation? (§D5: 'дельта качества неотличима от нуля')."""
    rows = []
    for f in sorted((COLD / "plain").glob("*.raw.txt")):
        cid = f.stem.replace(".raw", "")
        plain = f.read_text(encoding="utf-8")
        bfile = COLD / "banknote" / f"{cid}.clean.txt"
        if not bfile.exists():
            continue
        bank = bfile.read_text(encoding="utf-8")
        rows.append(dict(id=cid, plain_len=len(plain), banknote_clean_len=len(bank),
                         ratio=round(len(bank) / len(plain), 3) if plain else None))
    if rows:
        import statistics
        ratios = [r["ratio"] for r in rows if r["ratio"]]
        return rows, (round(statistics.mean(ratios), 3), round(statistics.pstdev(ratios), 3))
    return rows, (None, None)


def main():
    gt = X.load_gt()
    inj = injected_chunks_for_cold()
    cold = load_coldstart_chunks("plain")
    print(f"cold-start chunks: {len(cold)} (chapters {COLD_CHAPTERS}); injected comparison chunks: {len(inj)}\n")
    if not cold:
        print("no cold-start drafts yet — run coldstart_gen.py first"); return

    # 1. spread (confound removed)
    srows, _, _ = spread_compare(gt, inj, cold)
    print("== SPREAD: injected (confounded ~0) vs cold-start (real dispersion) ==")
    import statistics
    si = [r["spread_injected"] for r in srows]
    sc = [r["spread_coldstart"] for r in srows]
    print(f"  mean spread injected={statistics.mean(si):.3f} vs cold-start={statistics.mean(sc):.3f} (n={len(srows)} GT terms occ>=2)")
    for r in sorted(srows, key=lambda r: -r["spread_coldstart"])[:12]:
        print(f"    {r['src']:<6}({r['typ']:<7}) occ={r['occ']:<3} inj={r['spread_injected']} cold={r['spread_coldstart']}")

    # 2. canon-recovery
    crates, rows_cold = canon_compare(gt, inj, cold)
    print(f"\n== CANON-RECOVERY (§D5 hypothesis-1 gate: >=70% on cold-start) ==")
    print(f"  injected (confounded): {crates['injected'][0]:.3f} (n={crates['injected'][1]})")
    print(f"  cold-start (ECOLOGICAL): {crates['coldstart'][0]:.3f} (n={crates['coldstart'][1]})")

    # 3. A4/A5 banknote
    bank, bproposed = banknote_analysis(gt)
    print(f"\n== A4/A5 BANKNOTE CHANNEL ==")
    print(f"  entries={bank['n_entries']} distinct_src={bank['distinct_src']} banknote_lines={bank['n_banknote_lines']}")
    print(f"  parse_fail_chunks={bank['parse_fail_chunks']}/{bank['nchunks']} (rate {bank['parse_fail_rate']:.3f}) truncated={bank['truncated_chunks']}")
    print(f"  banknote recall: all={bank['recall_all']} f>=10={bank['recall_f_ge10']} f3-9={bank['recall_f3_9']} f<3={bank['recall_f_lt3']}")

    # marginal recall over best $0-arm (V-C) on f<3
    ar = A.Arms(X.load_chunks(), X.Contrast())
    a3set = A.candidate_set(ar.arm_A3())
    def marg_f_lt3():
        gained = []
        for e in gt:
            if X.freq_stratum(X.gt_occurrences(e, X.load_chunks())) != "f<3":
                continue
            in_a3 = any(s in a3set for s in e.norm_surfaces)
            in_bank = any(s in bproposed for s in e.norm_surfaces)
            if in_bank and not in_a3:
                gained.append(e.src)
        return gained
    gained = marg_f_lt3()
    print(f"  MARGINAL f<3 recall over V-C: banknote adds {gained} (V-C already had the rest)")

    # 4. quality delta
    qrows, (qmean, qstd) = quality_delta()
    print(f"\n== QUALITY DELTA (plain vs banknote clean draft) ==")
    print(f"  clean-length ratio banknote/plain: mean={qmean} std={qstd} (want ~1.0 = no degradation)")

    # cost
    ledger = BOOK / "exp16" / "coldstart_costs.jsonl"
    total = sum(json.loads(l).get("cost", 0) for l in open(ledger)) if ledger.exists() else 0
    print(f"\n== COST: ${total:.4f} (cap $10) ==")

    json.dump({"spread": srows, "spread_mean": {"injected": statistics.mean(si), "coldstart": statistics.mean(sc)},
               "canon": crates, "banknote": {k: v for k, v in bank.items() if k != "sample_entries"},
               "banknote_sample": bank["sample_entries"], "marginal_f_lt3": gained,
               "quality_delta": {"mean": qmean, "std": qstd}, "cost": total},
              open(BOOK / "exp16" / "coldstart_analysis.json", "w"), ensure_ascii=False, indent=1)
    print(f"[written] {BOOK/'exp16'/'coldstart_analysis.json'}")


if __name__ == "__main__":
    main()
