# Study S2C pre-specified tests (plan section 6): voice (Fisher exact), language (exact McNemar), cross-engine
# variance (chi-square homogeneity on S2's U and E answers), Holm-Bonferroni across the nine tests. Standard library
# only, so it runs with python -I.
#   python -I docs/evidence/research-programme-2026-q4/s2c_tests.py <s2c codes.json> <s2 codes.json>
#   python -I docs/evidence/research-programme-2026-q4/s2c_tests.py --selftest
import json
import math
import sys
from collections import defaultdict

ENGINES = ["openai", "anthropic", "google", "perplexity"]
ALPHA = 0.05


def fisher_exact_two_sided(a, b, c, d):
    """Two-sided Fisher exact p for the 2x2 table [[a, b], [c, d]]: sum of all tables with the same margins whose
    probability is <= the observed table's (the standard definition)."""
    r1, r2, c1, n = a + b, c + d, a + c, a + b + c + d
    def p_table(x):
        return math.comb(r1, x) * math.comb(r2, c1 - x) / math.comb(n, c1)
    p_obs = p_table(a)
    lo, hi = max(0, c1 - r2), min(r1, c1)
    return min(1.0, sum(p_table(x) for x in range(lo, hi + 1) if p_table(x) <= p_obs * (1 + 1e-9)))


def mcnemar_exact(b, c):
    """Exact two-sided McNemar p from the discordant counts b and c (binomial test, p = 0.5)."""
    n = b + c
    if n == 0:
        return 1.0
    k = min(b, c)
    return min(1.0, 2 * sum(math.comb(n, i) for i in range(0, k + 1)) / 2 ** n)


def chi2_sf(x, df):
    """Survival function of the chi-square distribution for integer df (regularised upper incomplete gamma)."""
    if x <= 0:
        return 1.0
    if df % 2 == 0:
        term, total = 1.0, 1.0
        for i in range(1, df // 2):
            term *= (x / 2) / i
            total += term
        return math.exp(-x / 2) * total
    # odd df: use the series through erfc
    s = math.erfc(math.sqrt(x / 2))
    term = math.sqrt(2 * x / math.pi) * math.exp(-x / 2)
    for i in range(1, (df - 1) // 2 + 1):
        s += term
        term *= x / (2 * i + 1)
    return min(1.0, s)


def chi2_homogeneity(rows):
    """rows: list of (successes, n). Returns (statistic, df, p) for a k x 2 test of equal proportions."""
    rows = [(k, n) for k, n in rows if n > 0]
    tot_k, tot_n = sum(k for k, _ in rows), sum(n for _, n in rows)
    if tot_n == 0 or tot_k in (0, tot_n) or len(rows) < 2:
        return 0.0, max(0, len(rows) - 1), 1.0
    p = tot_k / tot_n
    stat = sum((k - n * p) ** 2 / (n * p) + ((n - k) - n * (1 - p)) ** 2 / (n * (1 - p)) for k, n in rows)
    df = len(rows) - 1
    return stat, df, chi2_sf(stat, df)


def holm(pvals):
    """Holm-Bonferroni: list of (name, p) -> list of (name, p, adjusted p, rejected) in the input order."""
    order = sorted(range(len(pvals)), key=lambda i: pvals[i][1])
    m = len(pvals)
    adj = [0.0] * m
    running = 0.0
    for rank, i in enumerate(order):
        running = max(running, min(1.0, (m - rank) * pvals[i][1]))
        adj[i] = running
    return [(pvals[i][0], pvals[i][1], adj[i], adj[i] <= ALPHA) for i in range(m)]


def coded(doc):
    return [r for r in doc["answers"] if r.get("mode", "collect") == "collect" and r.get("code") is not None]


def run(s2c, s2):
    tests = []
    a = coded(s2c)
    # Voice: physician (VU3, VE3) against patient (VU1, VE1), F on eligibility.
    for e in ENGINES:
        phys = [r for r in a if r["engine"] == e and r["wordingId"] in ("VU3", "VE3")]
        pat = [r for r in a if r["engine"] == e and r["wordingId"] in ("VU1", "VE1")]
        fp, fq = sum(r["code"] == "F" for r in phys), sum(r["code"] == "F" for r in pat)
        p = fisher_exact_two_sided(fp, len(phys) - fp, fq, len(pat) - fq) if phys and pat else 1.0
        tests.append((f"voice {e}: physician {fp}/{len(phys)} v patient {fq}/{len(pat)}", p))
    # Language: p1 agreement (C) Arabic v English, paired by core/voice (wording pair LnE/LnA) and day/round.
    for e in ENGINES:
        idx = {(r["wordingId"], r.get("day"), r.get("round")): r for r in a if r["engine"] == e and r["family"] == "L"}
        b = c = 0
        for (wid, d, rd), ra in idx.items():
            if not wid.endswith("A"):
                continue
            re_ = idx.get((wid[:-1] + "E", d, rd))
            if re_ is None or ra.get("p1") is None or re_.get("p1") is None:
                continue
            ar_ok, en_ok = ra["p1"] == "C", re_["p1"] == "C"
            if ar_ok and not en_ok:
                b += 1
            elif en_ok and not ar_ok:
                c += 1
        tests.append((f"language {e}: discordant Arabic-only correct {b}, English-only correct {c}", mcnemar_exact(b, c)))
    # Cross-engine variance on S2's U and E families.
    s2rows = [r for r in coded(s2) if r["family"] in ("U", "E")]
    rows = [(sum(r["code"] == "F" for r in s2rows if r["engine"] == e), sum(1 for r in s2rows if r["engine"] == e)) for e in ENGINES]
    stat, df, p = chi2_homogeneity(rows)
    tests.append((f"cross-engine S2 U+E F rates {rows}: chi2 {stat:.2f} df {df}", p))
    for name, p, padj, rej in holm(tests):
        print(f"{'REJECT' if rej else 'retain'}  p={p:.4g}  Holm-adjusted={padj:.4g}  {name}")


def selftest():
    # Fisher: the classic tea-tasting table [[3,1],[1,3]] has two-sided p = 0.4857.
    assert abs(fisher_exact_two_sided(3, 1, 1, 3) - 0.4857) < 1e-3
    # McNemar exact: b=0, c=6 -> 2 * 0.5^6 = 0.03125.
    assert abs(mcnemar_exact(0, 6) - 0.03125) < 1e-9
    # Chi-square survival: P(chi2_3 > 7.815) ~ 0.05; P(chi2_2 > 5.991) ~ 0.05.
    assert abs(chi2_sf(7.815, 3) - 0.05) < 1e-3 and abs(chi2_sf(5.991, 2) - 0.05) < 1e-3
    # Homogeneity: equal proportions -> p = 1.
    assert chi2_homogeneity([(5, 10), (5, 10)])[2] > 0.99
    # Holm: p = (0.01, 0.04, 0.03) -> adjusted (0.03, 0.06, 0.06); only the first rejected at 0.05.
    res = holm([("a", 0.01), ("b", 0.04), ("c", 0.03)])
    assert [round(x[2], 4) for x in res] == [0.03, 0.06, 0.06] and [x[3] for x in res] == [True, False, False], res
    print("selftest: all checks pass")


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