# Lilly study 1 (orforglipron, UAE): reproducible counts from the committed codes file only.
#   python -I docs/evidence/lilly-orforglipron-uae/tally.py docs/evidence/lilly-orforglipron-uae/codes.json
# Rules: docs/strategy/lilly-orforglipron-uae-plan.md, sections 6 and 7. Repeated draws of one wording are not
# independent: rates carry Wilson 95% intervals for orientation, with per-wording code strings beside them, and no
# significance test is run. Sections marked "added, not pre-specified" are not in the plan.
import json
import math
import sys
from collections import Counter, defaultdict

ENGINES = ["openai", "anthropic", "google", "perplexity"]
STATUS_SET = ["D1", "D2", "D3", "P1", "P2", "K1", "B1", "B2"]
NAME_GROUPS = {"brand only": ["D1", "P1", "B1"], "INN only": ["D2", "P2"], "both": ["K1", "B2"], "descriptive": ["D3"]}
FAMILIES = {"D direct": ["D1", "D2", "D3"], "P patient": ["P1", "P2"], "K clinician": ["K1"], "B buyer": ["B1", "B2"]}


def wilson(k: int, n: int) -> str:
    if n == 0:
        return "n/a"
    z = 1.96
    p = k / n
    den = 1 + z * z / n
    centre = (p + z * z / (2 * n)) / den
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / den
    return f"{100 * max(0.0, centre - half):.0f}-{100 * min(1.0, centre + half):.0f}%"


def counts(rows):
    c = Counter(r["code"] for r in rows)
    return {k: c.get(k, 0) for k in ("A", "C", "B", "H", "0")}


def line(label, rows):
    n = len(rows)
    c = counts(rows)
    return (f"| {label} | {n} | {c['C']} ({wilson(c['C'], n)}) | {c['A']} ({wilson(c['A'], n)}) | {c['B']} | {c['H']} | {c['0']} |")


HEADER = "| Group | N | C correct (95% CI) | A false (95% CI) | B avail. neg. | H hedged | 0 not addressed |\n|---|---|---|---|---|---|---|"


def reading_a1(rows):
    c = counts(rows)
    if c["A"] == 0 and c["C"] >= 18:
        return "Current"
    if c["A"] >= 6:
        return "Not current"
    return "Mixed"


def has(r, flag):
    return any(f == flag or f.startswith(flag + ":") for f in r["flags"])


def main(path: str) -> None:
    sys.stdout.reconfigure(encoding="utf-8")
    doc = json.load(open(path, encoding="utf-8"))
    rows = doc["answers"]
    rounds = sorted({r["round"] for r in rows})
    print(f"answers coded: {len(rows)}; uncoded: {len(doc['uncoded'])}; rounds: {rounds}")
    by_engine = {e: [r for r in rows if r["engine"] == e] for e in ENGINES}

    print("\nA1 (primary). Foundayo status set: D1-D3, P1, P2, K1, B1, B2; 8 wordings x rounds, per engine")
    print(HEADER)
    for e in ENGINES:
        s = [r for r in by_engine[e] if r["wordingId"] in STATUS_SET]
        print(line(f"{e} ({reading_a1(s) if len(s) == 24 else 'incomplete'})", s))
    allst = [r for r in rows if r["wordingId"] in STATUS_SET]
    print(line("all four engines", allst))

    print("\nA2 (control). Mounjaro, M1-M3, per engine (control valid if C >= 7 of 9)")
    print(HEADER)
    for e in ENGINES:
        m = [r for r in by_engine[e] if r["family"] == "M"]
        valid = "valid" if len(m) == 9 and counts(m)["C"] >= 7 else ("not valid" if len(m) == 9 else "incomplete")
        print(line(f"{e} (control {valid})", m))

    print("\nA3 (currency gap). C rate on Mounjaro minus C rate on Foundayo status set, percentage points")
    for e in ENGINES:
        m = [r for r in by_engine[e] if r["family"] == "M"]
        s = [r for r in by_engine[e] if r["wordingId"] in STATUS_SET]
        if m and s:
            gap = 100 * counts(m)["C"] / len(m) - 100 * counts(s)["C"] / len(s)
            print(f"  {e}: Mounjaro {counts(m)['C']}/{len(m)}, Foundayo {counts(s)['C']}/{len(s)}, gap {gap:+.0f} points")

    for title, groups in (("A4 (name)", NAME_GROUPS), ("A5 (frame)", FAMILIES)):
        print(f"\n{title}. Foundayo status set by group, pooled over engines")
        print(HEADER)
        for g, ids in groups.items():
            print(line(g, [r for r in rows if r["wordingId"] in ids]))
        print("  A by group and engine:", "; ".join(
            f"{g}: " + ", ".join(f"{e} {counts([r for r in by_engine[e] if r['wordingId'] in ids])['A']}/{len([r for r in by_engine[e] if r['wordingId'] in ids])}" for e in ENGINES)
            for g, ids in groups.items()))

    print("\nA6 (search). Foundayo status set")
    for e in ENGINES:
        s = [r for r in by_engine[e] if r["wordingId"] in STATUS_SET]
        a = [r for r in s if r["code"] == "A"]
        c = [r for r in s if r["code"] == "C"]
        print(f"  {e}: searched {sum(r['searched'] for r in s)}/{len(s)}; A answers with no search {sum(not r['searched'] for r in a)}/{len(a)}; "
              f"C answers that searched {sum(r['searched'] for r in c)}/{len(c)}; mean searches {sum(r['searchCount'] for r in s) / max(1, len(s)):.1f}")

    print("\nA7 (discovery). Family O (O1, O2): answers naming Foundayo or orforglipron, and their codes")
    for e in ENGINES:
        o = [r for r in by_engine[e] if r["family"] == "O"]
        named = [r for r in o if "unnamed" not in r["flags"] and r["code"] != "0"]
        print(f"  {e}: named {len(named)}/{len(o)}; codes {dict(Counter(r['code'] for r in o))}")

    print("\nA8 (secondary). All 30 Foundayo answers per engine (status set and O); Mounjaro where stated")
    for e in ENGINES:
        f = [r for r in by_engine[e] if r["subject"] == "Foundayo"]
        m = [r for r in by_engine[e] if r["subject"] == "Mounjaro"]
        ind_f = Counter(fl for r in f for fl in r["flags"] if fl.startswith("ind:"))
        ind_m = Counter(fl for r in m for fl in r["flags"] if fl.startswith("ind:"))
        print(f"  {e}: stale {sum(has(r, 'stale') for r in f)}, fda {sum(has(r, 'fda') for r in f)}, euWrong {sum(has(r, 'euWrong') for r in f)}, "
              f"noRec {sum(has(r, 'noRec') for r in f)}, W {sum(has(r, 'W') for r in f)}, contradictory {sum(has(r, 'contradictory') for r in f)}, "
              f"availNeg {sum(has(r, 'availNeg') for r in f)}; form right {sum(has(r, 'formRight') for r in f)}/{len(f)}, wrong {sum(has(r, 'formWrong') for r in f)}; "
              f"Foundayo UAE indication stated {dict(sorted(ind_f.items()))} of {len(f)}; Mounjaro UAE indication stated {dict(sorted(ind_m.items()))} of {len(m)}; "
              f"truncated {sum(bool(r['truncated']) for r in by_engine[e])}/{len(by_engine[e])}")
    b = [r for r in by_engine["openai"] if r["wordingId"] in ("B1", "B2")]
    print(f"  B1 and B2 on gpt-4o-mini: {dict(Counter(r['code'] for r in b))} of {len(b)} (Field Note 07 Stage 1, 13-16 Oct: C 7, H 1 of 8)")

    print("\nBy wording (codes in round order), per engine")
    for e in ENGINES:
        by = defaultdict(list)
        for r in sorted(by_engine[e], key=lambda r: (r["wordingId"], r["round"])):
            by[r["wordingId"]].append(r["code"])
        print(f"  {e}: " + "  ".join(f"{w} {''.join(v)}" for w, v in sorted(by.items())))

    print("\nCited page classes (pages per class, summed over answers); Gemini by class only")
    for e in ENGINES:
        c = Counter()
        for r in by_engine[e]:
            c.update(r["citedHostClasses"])
        reg = sum(1 for r in by_engine[e] if r["citedHostClasses"].get("regulatorUAE"))
        print(f"  {e}: {dict(sorted(c.items()))}; answers citing a UAE regulator page {reg}/{len(by_engine[e])}")

    print("\nKeyword cross-check: answers whose keyword code differs from the hand code")
    dis = [r for r in rows if r["keywordDisagrees"]]
    print(f"  {len(dis)} of {len(rows)}: " + ", ".join(f"r{r['round']} {r['engine']} {r['wordingId']} (hand {r['code']}, keyword {r['keywordCode']})" for r in dis))

    print("\nAdded, not pre-specified: what each answer says Mounjaro is approved for in the UAE, by engine and wording")
    print("  (from the ind: flag; '-' = no UAE approved-use statement; no UAE label is in hand, so none is coded right or wrong)")
    for e in ENGINES:
        cells = []
        for w in ("M1", "M2", "M3"):
            rs = sorted((r for r in by_engine[e] if r["wordingId"] == w), key=lambda r: r["round"])
            cells.append(f"{w} " + " ".join(next((f[4:] for f in r["flags"] if f.startswith("ind:")), "-") for r in rs))
        print(f"  {e}: " + "; ".join(cells))
    m_all = [r for r in rows if r["subject"] == "Mounjaro"]
    ind_all = Counter(next((f[4:] for f in r["flags"] if f.startswith("ind:")), "-") for r in m_all)
    print(f"  pooled over {len(m_all)} Mounjaro answers: {dict(sorted(ind_all.items()))}")
    print("  (BOTH = type 2 diabetes and weight management; T2Dmain = type 2 diabetes as the main approved use, weight")
    print("   management not called off-label; T2Donly = type 2 diabetes only, weight management called off-label)")

    def readings(e):
        return {next((f[4:] for f in r["flags"] if f.startswith("ind:")), None) for r in by_engine[e] if r["subject"] == "Mounjaro"}

    split = [e for e in ENGINES if "BOTH" in readings(e) and readings(e) & {"T2Dmain", "T2Donly"}]
    print(f"  engines giving BOTH and a type 2 diabetes reading across their own Mounjaro answers: {split}")
    print(f"  BOTH answers that also state the other reading (indContradictory): {sum(has(r, 'indContradictory') for r in m_all)}")
    for lab in ("BOTH", "T2Dmain", "T2Donly"):
        g = [r for r in m_all if f"ind:{lab}" in r["flags"]]
        print(f"  {lab} answers citing at least one UAE clinic, pharmacy or telehealth page: "
              f"{sum(1 for r in g if r['citedHostClasses'].get('uaeCareOrPharmacy'))}/{len(g)}")
    f_all = [r for r in rows if r["subject"] == "Foundayo"]
    print(f"  Foundayo answers stating a UAE approved use: {sum(1 for r in f_all if any(f.startswith('ind:') for f in r['flags']))}; "
          f"of them weight management: {sum(1 for r in f_all if 'ind:WM' in r['flags'])}")
    print("  BOTH answers by wording: " + ", ".join(f"{w} {sum(1 for r in m_all if r['wordingId'] == w and 'ind:BOTH' in r['flags'])}" for w in ("M1", "M2", "M3")))

    print("\nAdded, not pre-specified: other observations flagged during coding")
    if "note" not in rows[0]:
        print("  coder notes are withheld from this file; flagged answers only:")
    for r in rows:
        note = r.get("note") or ""
        if any(f.startswith(("added:", "availPos")) for f in r["flags"]) or note.startswith(("ADDED OBSERVATION", "DISPUTED READING", "RECODE")):
            print(f"  r{r['round']} {r['engine']} {r['wordingId']} ({r['responseIdTail']}) code {r['code']} flags {r['flags']}" + (f": {note}" if note else ""))

    print("\nAdded, not pre-specified: sources in the Foundayo status set (24 answers per engine; Gemini by grounding-source class)")
    for e in ENGINES:
        st = [r for r in by_engine[e] if r["wordingId"] in STATUS_SET]
        c = Counter()
        for r in st:
            c.update(r["citedHostClasses"])
        print(f"  {e}: answers with a UAE regulator page among their sources {sum(1 for r in st if r['citedHostClasses'].get('regulatorUAE'))}/{len(st)}; "
              f"pages by class {dict(sorted(c.items()))}")
    pooled = Counter()
    for r in rows:
        if r["wordingId"] in STATUS_SET:
            pooled.update(r["citedHostClasses"])
    print(f"  all four, pages by class over the 96 status answers ({sum(pooled.values())} pages): {dict(sorted(pooled.items()))}")
    print("  answers citing glpguidebook.com (the site discussed in Field Note 07): "
          + ", ".join(f"r{r['round']} {r['engine']} {r['wordingId']}" for r in rows if r.get("citesFieldNote07Site")))

    print("\nAdded, not pre-specified: gpt-4o-mini search queries containing '2023'")
    o = by_engine["openai"]
    print(f"  {sum(r['queryHas2023'] for r in o)} of {len(o)} answers; their codes {dict(Counter(r['code'] for r in o if r['queryHas2023']))}")


if __name__ == "__main__":
    main(sys.argv[1])