#!/usr/bin/env python3
"""exp16 — arm protocol driver (research/20 §D5). $0.
Phases:
  tune  — ch1-15 (tuning half): threshold sensitivity (±50% on floor/over_rep/subsume/λ), pick frozen config.
  test  — ch16-25 (test half): frozen config, PRIMARY reported result (anti-overfit §D1).
  full  — all 57 chunks: catastrophe screen + trade-off curve + candidate dump for canon/alias/precision@30.
Usage: python run_arms.py {tune|test|full}
"""
from __future__ import annotations
import json
import sys
from pathlib import Path

import exp16_common as X
import arms as A

OUT = Path("/home/ubuntu/books/gu-zhenren/exp16")
OUT.mkdir(parents=True, exist_ok=True)


def build(cfg, chunks, contrast, lam_for_arms):
    ar = A.Arms(chunks, contrast, cfg=cfg)
    return {"A1": ar.arm_A1(), "A2": ar.arm_A2(lam=lam_for_arms),
            "A3": ar.arm_A3(lam=0.0), "A3-abl": ar.arm_A3_abl()}, ar


def recall_f3(ranked, gt, chunks):
    rp, _ = A.recall_by(ranked, gt, chunks, top_k=None)
    # f>=3 = union of f>=10 and f3-9
    hits = tot = 0
    for e in gt:
        f = X.gt_occurrences(e, chunks)
        if f >= 3:
            tot += 1
            cset = A.candidate_set(ranked)
            hits += any(s in cset for s in e.norm_surfaces)
    return hits / tot if tot else None, rp


def phase_tune(chunks, gt, contrast):
    tune_ch = X.TUNING_CH
    sub = [c for c in chunks if c.chapter in tune_ch]
    gt_t = A.gt_in_chapters(gt, chunks, tune_ch)
    print(f"== TUNING half (ch1-15): {len(sub)} chunks, {len(gt_t)} GT terms occurring here ==\n")
    base = dict(A.FROZEN)
    results = {}

    def run(cfg, tag):
        armset, ar = build(cfg, sub, contrast, lam_for_arms=0.0)
        rA1_f3, rpA1 = recall_f3(armset["A1"], gt_t, sub)
        rA3_f3, rpA3 = recall_f3(armset["A3"], gt_t, sub)
        ok3, _ = A.catastrophe_screen(armset["A3"], 50)
        results[tag] = dict(A1_recall_f3=round(rA1_f3, 3), A3_recall_f3=round(rA3_f3, 3),
                            A3_recall_f_lt3=rpA3["f<3"][0] and round(rpA3["f<3"][0], 3),
                            A3_recall_overall=round(rpA3["overall"][0], 3),
                            catastrophe=ok3, n_A3=len(armset["A3"]),
                            pseudo_prec_A3=round(A.pseudo_precision(armset["A3"], gt_t, cfg["top_k"]), 3))
        r = results[tag]
        print(f"  {tag:<26} A1@f>=3={r['A1_recall_f3']} A3@f>=3={r['A3_recall_f3']} "
              f"A3@f<3={r['A3_recall_f_lt3']} A3@all={r['A3_recall_overall']} "
              f"cat={'ok' if r['catastrophe'] else 'FAIL'} nA3={r['n_A3']} pp={r['pseudo_prec_A3']}")

    print("baseline (frozen candidate defaults):")
    run(base, "floor3_over15_sub0.80")
    print("\nsensitivity — freq_floor (±50%: 2,3,4):")
    for fl in (2, 3, 4):
        c = dict(base); c["freq_floor"] = fl
        run(c, f"floor{fl}")
    print("\nsensitivity — formant_min_over_rep (±50%: 7.5,15,22.5,30):")
    for orp in (7.5, 15.0, 22.5, 30.0):
        c = dict(base); c["formant_min_over_rep"] = orp
        run(c, f"over_rep{orp}")
    print("\nsensitivity — subsume_alpha (0.5,0.8,1.0[off]):")
    for sa in (0.5, 0.8, 1.0):
        c = dict(base); c["subsume_alpha"] = sa
        run(c, f"subsume{sa}")
    print("\nsensitivity — formant_min_partners (2,3,4):")
    for mp in (2, 3, 4):
        c = dict(base); c["formant_min_partners"] = mp
        run(c, f"partners{mp}")

    json.dump(results, open(OUT / "tuning_sensitivity.json", "w"), ensure_ascii=False, indent=1)
    print(f"\n[written] {OUT/'tuning_sensitivity.json'}")
    print("\nFROZEN choice: freq_floor=3, formant_min_over_rep=15, subsume_alpha=0.80, min_partners=3, lam=0 "
          "(injected drafts) — stable across the sweep; see report §1.")


def phase_test(chunks, gt, contrast):
    test_ch = X.TEST_CH
    sub = [c for c in chunks if c.chapter in test_ch]
    gt_t = A.gt_in_chapters(gt, chunks, test_ch)
    print(f"== TEST half (ch16-25): {len(sub)} chunks, {len(gt_t)} GT terms occurring here — FROZEN config ==\n")
    armset, ar = build(dict(A.FROZEN), sub, contrast, lam_for_arms=A.FROZEN["lam"])
    K = A.FROZEN["top_k"]
    out = {}
    for name, ranked in armset.items():
        out[name] = A.summarize(name, ranked, gt_t, sub, K)
        s = out[name]; rp = s["recall_proposed"]
        print(f"{name}: catastrophe={'PASS' if s['catastrophe_ok'] else 'FAIL'}")
        print(f"   recall@PROPOSED overall={rp['overall'][0]} f>=10={rp['f>=10'][0]} f3-9={rp['f3-9'][0]} f<3={rp['f<3'][0]}")
        print(f"   by type: " + " ".join(f"{t}={rp[t][0]}({rp[t][1]})" for t in ("name","title","place","term","nickname") if rp[t][0] is not None))
        print(f"   misses@proposed: {s['misses_proposed']}")
    json.dump(out, open(OUT / "test_half_result.json", "w"), ensure_ascii=False, indent=1)
    print(f"\n[written] {OUT/'test_half_result.json'}")


def phase_full(chunks, gt, contrast):
    print(f"== FULL slice (57 chunks) — FROZEN config; catastrophe + trade-off + candidate dump ==\n")
    armset, ar = build(dict(A.FROZEN), chunks, contrast, lam_for_arms=A.FROZEN["lam"])
    K = A.FROZEN["top_k"]
    out = {"formants": {c: ar.formants[c] for c in ar.formants}}
    for name, ranked in armset.items():
        out[name] = A.summarize(name, ranked, gt, chunks, K)
    # trade-off curve for the winner A3 (recall@K vs pseudo-precision@K)
    A3 = armset["A3"]
    curve = []
    for k in (30, 50, 90, 150, 250, 400, 700, len(A3)):
        rk, _ = A.recall_by(A3, gt, chunks, top_k=k)
        curve.append(dict(k=k, recall_overall=round(rk["overall"][0], 3),
                          recall_f3=round((sum(1 for e in gt if X.gt_occurrences(e,chunks)>=3 and
                              any(({c.src:i for i,c in enumerate(A3)}.get(s,10**9))<k for s in e.norm_surfaces))
                              / sum(1 for e in gt if X.gt_occurrences(e,chunks)>=3)), 3),
                          pseudo_prec=round(A.pseudo_precision(A3, gt, k), 3)))
    out["A3_tradeoff_curve"] = curve
    # dump A3 candidates (for canon/alias/precision@30 stages) + spread/dst for freq core
    dump = []
    for c in armset["A3"]:
        dump.append(dict(src=c.src, score=round(c.score, 2), freq=c.freq, types=c.types,
                         evidence=c.evidence[:4], dst_variants=c.dst_variants[:4],
                         spread=round(c.spread, 3), from_pattern=c.from_pattern))
    json.dump(dump, open(OUT / "A3_candidates_full.json", "w"), ensure_ascii=False, indent=1)
    json.dump(out, open(OUT / "full_slice_result.json", "w"), ensure_ascii=False, indent=1)
    print("A3 trade-off (recall vs pseudo-precision):")
    for row in curve:
        print(f"   top-{row['k']:<5} recall_all={row['recall_overall']} recall_f>=3={row['recall_f3']} pseudo-prec={row['pseudo_prec']}")
    print(f"\n[written] {OUT/'full_slice_result.json'}, {OUT/'A3_candidates_full.json'} ({len(dump)} candidates)")


if __name__ == "__main__":
    phase = sys.argv[1] if len(sys.argv) > 1 else "tune"
    chunks = X.load_chunks(); gt = X.load_gt(); contrast = X.Contrast()
    {"tune": phase_tune, "test": phase_test, "full": phase_full}[phase](chunks, gt, contrast)