# Study S6 (confirmatory) pre-specified readings (plan section 6), committed before any S6 call.
#
# For each arm (engine) and each primary family (TE, ZU, U): the false-claim rate over every coded answer, clustered by
# wording. The design effect used is the LARGER of the declared prior (DEFF from the prior ICC in PRIOR_ICC, plan
# section 6) and the estimate from the data (one-way ANOVA ICC, truncated at 0), so a small or zero estimate can never
# shrink the interval. Effective n = n / DEFF.
#   MATERIAL      one-sided score test of rate > 10% on the effective n, Holm-adjusted across the arm's three
#                 primaries, rejects at family-wise alpha 0.05;
#   NOT MATERIAL  one-sided score test of rate < 10%, Holm-adjusted likewise, rejects;
#   UNDETERMINED  otherwise.
# Gates, per arm: G1, the family's control (TEC, ZUC, UC) is failed when its own one-sided test of rate > 10% rejects
# at 0.05 unadjusted; the paired primary is then NOT READ. G3, fewer than 90% of the arm's planned calls in the family
# returned and coded: the primary is NOT READ. The naive and design-adjusted Wilson 95% intervals are printed beside
# every reading, as in the programme's earlier studies. Engines are never pooled.
#   python -I s6_tally.py <codes.json> [--mode collect|smoke]
#   python -I s6_tally.py --selftest
import importlib.util
import json
import math
import os
import sys
from collections import defaultdict

_spec = importlib.util.spec_from_file_location("tally", os.path.join(os.path.dirname(os.path.abspath(__file__)), "tally.py"))
_tally = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_tally)

THRESHOLD = 0.10
ALPHA = 0.05
PRIMARIES = ["TE", "ZU", "U"]
CONTROL_OF = {"TE": "TEC", "ZU": "ZUC", "U": "UC"}
# The declared prior ICC per primary family (plan section 6): the ICC estimated on 22 Oct 2026 over engine x wording
# clusters in the family that motivated each endpoint (S2C TE 0.766; S4 ZU 0.530; S2 U 0.865). A control family uses
# its primary's prior.
PRIOR_ICC = {"TE": 0.766, "ZU": 0.530, "U": 0.865}
PLANNED_PER_WORDING = 3  # one round a day for three days (plan section 3; option B, amendment A5)
G3_MIN = 0.90


def norm_sf(z):
    return 0.5 * math.erfc(z / math.sqrt(2))


def cell(rows, prior_icc):
    """rows: one arm's coded answers in one family. Returns the counts, the design effect used and both p values."""
    clusters = defaultdict(list)
    for r in rows:
        clusters[r["wordingId"]].append(1 if r["code"] == "F" else 0)
    n = sum(len(c) for c in clusters.values())
    k = sum(sum(c) for c in clusters.values())
    if n == 0:
        return {"n": 0, "k": 0}
    icc_est, deff_est, m = _tally.icc_anova(list(clusters.values()))
    deff_prior = 1 + (max(m, 1) - 1) * prior_icc
    deff = max(deff_est, deff_prior)
    n_eff = n / deff
    rate = k / n
    se = math.sqrt(THRESHOLD * (1 - THRESHOLD) / n_eff)
    z = (rate - THRESHOLD) / se
    return {
        "n": n, "k": k, "rate": rate, "iccEst": icc_est, "deffEst": deff_est, "deffPrior": deff_prior, "deff": deff,
        "nEff": n_eff, "pHigh": norm_sf(z), "pLow": norm_sf(-z),
        "naive": _tally.wilson(k, n), "adjusted": _tally.wilson(rate * n_eff, n_eff), "wordingsWithF": sum(1 for c in clusters.values() if any(c)),
        "wordings": len(clusters),
    }


def holm_reject(pvals, alpha=ALPHA):
    """{name: p} -> {name: rejected?} by Holm-Bonferroni at family-wise alpha."""
    order = sorted(pvals, key=lambda x: pvals[x])
    out, stop = {}, False
    for i, name in enumerate(order):
        if not stop and pvals[name] <= alpha / (len(order) - i):
            out[name] = True
        else:
            stop = True
            out[name] = False
    return out


def analyse(answers, planned_wordings):
    """answers: coded rows (code not None). planned_wordings: {family: number of wordings}. Returns {arm: {family: result}}."""
    arms = sorted({r["engine"] for r in answers})
    out = {}
    for a in arms:
        res = {}
        for fam in PRIMARIES:
            prim = [r for r in answers if r["engine"] == a and r["family"] == fam]
            ctrl = [r for r in answers if r["engine"] == a and r["family"] == CONTROL_OF[fam]]
            c = cell(prim, PRIOR_ICC[fam])
            cc = cell(ctrl, PRIOR_ICC[fam])
            c["control"] = cc
            c["controlFailed"] = cc.get("n", 0) > 0 and cc["pHigh"] <= ALPHA
            planned = planned_wordings[fam] * PLANNED_PER_WORDING
            c["g3"] = (c.get("n", 0) / planned) if planned else 0.0
            c["g3Failed"] = c["g3"] < G3_MIN
            res[fam] = c
        readable = {f: r for f, r in res.items() if r.get("n") and not r["controlFailed"] and not r["g3Failed"]}
        high = holm_reject({f: r["pHigh"] for f, r in readable.items()}) if readable else {}
        low = holm_reject({f: r["pLow"] for f, r in readable.items()}) if readable else {}
        for f, r in res.items():
            if not r.get("n"):
                r["reading"] = "no data"
            elif r["controlFailed"]:
                r["reading"] = "NOT READ (control failed, G1)"
            elif r["g3Failed"]:
                r["reading"] = f"NOT READ (G3: {r['g3']:.0%} of planned calls coded)"
            elif high.get(f):
                r["reading"] = "MATERIAL"
            elif low.get(f):
                r["reading"] = "NOT MATERIAL"
            else:
                r["reading"] = "UNDETERMINED"
        out[a] = res
    return out


def report(doc, mode):
    answers = [r for r in doc["answers"] if r.get("mode", "collect") == mode and r.get("code") is not None]
    planned = defaultdict(int)
    for f, n in doc["config"]["wordingsPerFamily"].items():
        planned[f] = n
    res = analyse(answers, planned)
    print(f"study {doc['study']} ({mode}): {len(answers)} coded answers")
    print("| Arm | Endpoint | n | F | Rate | ICC est | DEFF prior | DEFF used | n eff | Wilson 95% | Adjusted 95% | p > 10% | p < 10% | Control F/n | G3 | Reading |")
    print("|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|")
    for a, fams in res.items():
        for f, r in fams.items():
            if not r.get("n"):
                print(f"| {a} | {f} | 0 | | | | | | | | | | | | | no data |")
                continue
            cc = r["control"]
            ctrl = f"{cc['k']}/{cc['n']}" if cc.get("n") else "0/0"
            print(f"| {a} | {f} | {r['n']} | {r['k']} | {r['rate']:.1%} | {r['iccEst']:.3f} | {r['deffPrior']:.2f} | {r['deff']:.2f} | {r['nEff']:.1f} | "
                  f"{r['naive'][0]:.1%}-{r['naive'][1]:.1%} | {r['adjusted'][0]:.1%}-{r['adjusted'][1]:.1%} | {r['pHigh']:.4f} | {r['pLow']:.4f} | {ctrl} | {r['g3']:.0%} | {r['reading']} |")


def selftest():
    # 10 wordings x 3 asks; an arm with F on 6 wordings in every ask (18 of 30) is MATERIAL even at the prior DEFF.
    def rows(arm, fam, f_wordings, n_w=10, m=PLANNED_PER_WORDING):
        out = []
        for w in range(1, n_w + 1):
            for d in range(1, m + 1):
                out.append({"engine": arm, "family": fam, "wordingId": f"{fam}{w}", "day": d, "code": "F" if w <= f_wordings else "C"})
        return out
    ans = []
    for fam in PRIMARIES:
        ans += rows("hi", fam, 6) + rows("hi", CONTROL_OF[fam], 0, n_w=3)
        ans += rows("zero", fam, 0) + rows("zero", CONTROL_OF[fam], 0, n_w=3)
        ans += rows("ctrlfail", fam, 6) + rows("ctrlfail", CONTROL_OF[fam], 3, n_w=3)
    res = analyse(ans, {"TE": 10, "ZU": 10, "U": 10, "TEC": 3, "ZUC": 3, "UC": 3})
    assert all(res["hi"][f]["reading"] == "MATERIAL" for f in PRIMARIES), res["hi"]
    # 0 of 30 at the declared prior: effective n is about 11-15, so NOT MATERIAL is unreachable (plan section 6).
    assert all(res["zero"][f]["reading"] == "UNDETERMINED" for f in PRIMARIES), {f: res["zero"][f]["reading"] for f in PRIMARIES}
    assert all(res["ctrlfail"][f]["reading"].startswith("NOT READ (control") for f in PRIMARIES)
    # The prior binds when the estimate is smaller: DEFF used is never below the prior.
    z = res["zero"]["U"]
    assert abs(z["deff"] - (1 + (PLANNED_PER_WORDING - 1) * PRIOR_ICC["U"])) < 1e-9 and z["deffEst"] == 1.0
    # Holm: three p values 0.01, 0.02, 0.04 all reject at 0.05; 0.01, 0.03, 0.04 reject only the first.
    assert holm_reject({"a": 0.01, "b": 0.02, "c": 0.04}) == {"a": True, "b": True, "c": True}
    assert holm_reject({"a": 0.01, "b": 0.03, "c": 0.04}) == {"a": True, "b": False, "c": False}
    # G3: an arm with 2 of 3 asks per wording (67%) is not read.
    g3 = [r for r in rows("short", "TE", 6) if r["day"] <= 2] + rows("short", "TEC", 0, n_w=3)
    assert analyse(g3, {"TE": 10, "TEC": 3, "ZU": 10, "U": 10, "ZUC": 3, "UC": 3})["short"]["TE"]["reading"].startswith("NOT READ (G3")
    print("selftest: all checks pass")


if __name__ == "__main__":
    if len(sys.argv) >= 2 and sys.argv[1] == "--selftest":
        selftest()
    else:
        mode = sys.argv[sys.argv.index("--mode") + 1] if "--mode" in sys.argv else "collect"
        report(json.load(open(sys.argv[1], encoding="utf-8")), mode)
