"""Serious adverse events in randomized placebo-controlled trials of GLP-1 receptor agonists.

Reads the ClinicalTrials.gov records in data/ctgov/, the screening and term coding decided
before any arm was compared (data/screening.json, data/term_coding.json), and the outcome
definitions (data/outcomes.json), and writes every pre-registered estimate under results/.

Usage: python3 code/analyze.py <bundle root> [--blind]
--blind swaps the arms of each comparison at random, as the code was developed and tested
before the plan was registered; its results say nothing about the drugs.
"""

import csv
import json
import math
import os
import sys

import numpy as np
from scipy import stats

ROOT = sys.argv[1] if len(sys.argv) > 1 else "."
BLIND = "--blind" in sys.argv
OUT = os.path.join(ROOT, "results")


def load(path):
    with open(os.path.join(ROOT, path), encoding="utf-8") as f:
        return json.load(f)


def norm(text):
    """A term or organ system as the coding keys it: spaces collapsed, lower case."""
    return " ".join((text or "").split()).lower()


P = load("data/parameters.json")
OUTCOMES = load("data/outcomes.json")
MOLECULES = load("data/molecules.json")
SCREENING = {row["nct"]: row for row in load("data/screening.json")}
CODING = {row["term"]: row for row in load("data/term_coding.json")}


# ---------------------------------------------------------------------------
# Comparisons: each a stratum with a GLP-1 arm and a placebo arm, pooled over groups.
# ---------------------------------------------------------------------------


class Arm:
    def __init__(self):
        self.at_risk = 0
        self.any_serious = 0
        self.deaths = 0
        self.deaths_at_risk = 0
        self.deaths_known = True
        self.terms = {}  # normalized term -> participants affected (summed over groups and rows)
        self.socs = {}  # normalized organ system -> {term -> participants affected}


def group_at_risk(group, rows):
    """Participants at risk of a serious event in a group: as the group reports it, else the most
    any serious term reports for it, else its other-event or death denominators."""
    if group.get("seriousNumAtRisk"):
        return group["seriousNumAtRisk"]
    at_risk = [s.get("numAtRisk") or 0 for e in rows for s in e.get("stats", []) if s.get("groupId") == group["id"]]
    return max(at_risk, default=0) or group.get("otherNumAtRisk") or group.get("deathsNumAtRisk") or 0


def build_arm(record, group_ids):
    ae = record["resultsSection"]["adverseEventsModule"]
    groups = {g["id"]: g for g in ae["eventGroups"]}
    rows = ae.get("seriousEvents", [])
    arm = Arm()
    for gid in group_ids:
        g = groups[gid]
        arm.at_risk += group_at_risk(g, rows)
        arm.any_serious += g.get("seriousNumAffected") or 0
        if g.get("deathsNumAffected") is None or not g.get("deathsNumAtRisk"):
            arm.deaths_known = False
        else:
            arm.deaths += g["deathsNumAffected"]
            arm.deaths_at_risk += g["deathsNumAtRisk"]
    wanted = set(group_ids)
    for e in rows:
        term, soc = norm(e.get("term")), norm(e.get("organSystem"))
        n = sum(s.get("numAffected") or 0 for s in e.get("stats", []) if s.get("groupId") in wanted)
        if n == 0:
            continue
        arm.terms[term] = arm.terms.get(term, 0) + n
        arm.socs.setdefault(soc, {})
        arm.socs[soc][term] = arm.socs[soc].get(term, 0) + n
    return arm


def build_strata():
    strata = []
    for nct in sorted(SCREENING):
        row = SCREENING[nct]
        if not row["eligible"]:
            continue
        with open(os.path.join(ROOT, "data", "ctgov", f"{nct}.json"), encoding="utf-8") as f:
            record = json.load(f)
        for k, comp in enumerate(row["comparisons"]):
            glp1 = build_arm(record, comp["glp1_groups"])
            placebo = build_arm(record, comp["placebo_groups"])
            if glp1.at_risk == 0 or placebo.at_risk == 0:
                continue
            classes = {MOLECULES[m] for m in comp["molecules"]}
            strata.append({
                "id": f"{nct}#{k + 1}",
                "nct": nct,
                "glp1": glp1,
                "placebo": placebo,
                "molecules": sorted(comp["molecules"]),
                "selective": classes == {"selective"},
                "population": row["population"],
            })
    if BLIND:
        rng = np.random.default_rng(P["blind_seed"])
        for s in strata:
            if rng.random() < 0.5:
                s["glp1"], s["placebo"] = s["placebo"], s["glp1"]
    return strata


# ---------------------------------------------------------------------------
# Counting an outcome in an arm
# ---------------------------------------------------------------------------

TERMS_OF = {}  # outcome id -> set of normalized terms coded to it
for term, row in CODING.items():
    for oid in row["outcomes"]:
        TERMS_OF.setdefault(oid, set()).add(term)

CANONICAL_TERMS = {}  # canonical preferred term -> set of normalized terms
for term, row in CODING.items():
    CANONICAL_TERMS.setdefault(row["canonical"], set()).add(term)


def by_preferred_term(arm, outcome):
    """An outcome's participants affected in an arm, by preferred term: {canonical: affected}.

    An outcome is an id from data/outcomes.json, ("soc", name), or ("pt", canonical term).
    """
    if outcome[0] == "soc":
        terms = arm.socs.get(outcome[1], {})
    elif outcome[0] == "pt":
        terms = {t: arm.terms.get(t, 0) for t in CANONICAL_TERMS[outcome[1]]}
    else:
        terms = {t: arm.terms.get(t, 0) for t in TERMS_OF.get(outcome, ())}
    out = {}
    for term, n in terms.items():
        if n:
            canonical = CODING[term]["canonical"] if term in CODING else term
            out[canonical] = out.get(canonical, 0) + n
    return out


def count(arm, outcome, only=None):
    """Participants with the outcome in an arm: (affected, at risk), or None if not reported.

    An outcome is "death", "any_serious", or anything `by_preferred_term` takes, counted as the
    sum over its terms of participants affected (an upper bound on participants with the
    outcome), capped at those at risk. `only` restricts it to one preferred term.
    """
    if outcome == "death":
        return (arm.deaths, arm.deaths_at_risk) if arm.deaths_known else None
    if outcome == "any_serious":
        return arm.any_serious, arm.at_risk
    counts = by_preferred_term(arm, outcome)
    affected = counts.get(only, 0) if only is not None else sum(counts.values())
    return min(affected, arm.at_risk), arm.at_risk


def tables(strata, outcome, only=None):
    """One 2x2 table per stratum: (a, n1, c, n0) = GLP-1 affected, at risk, placebo affected, at risk."""
    out = []
    for s in strata:
        g, p = count(s["glp1"], outcome, only), count(s["placebo"], outcome, only)
        if g is None or p is None:
            continue
        out.append((s, g[0], g[1], p[0], p[1]))
    return out


def pooled_by_preferred_term(strata, outcome):
    """For each stratum, an outcome's participants affected by preferred term, both arms together."""
    out = []
    for s in strata:
        both = by_preferred_term(s["glp1"], outcome)
        for term, n in by_preferred_term(s["placebo"], outcome).items():
            both[term] = both.get(term, 0) + n
        out.append(both)
    return out


def dominant_term(strata, outcome):
    """The outcome's preferred term with the most participants affected, both arms together."""
    totals = {}
    for both in pooled_by_preferred_term(strata, outcome):
        for term, n in both.items():
            totals[term] = totals.get(term, 0) + n
    return min(totals, key=lambda t: (-totals[t], t)) if totals else None


def overlap_bound(strata, outcome, overall):
    """How far an outcome's count could overstate participants, blind to arm: its count over all
    strata divided by the sum, over strata, of its largest preferred term (each stratum has at
    least that many participants with the outcome), and no more than `overall`, what the data
    show for serious events as a whole."""
    total = largest = 0
    for both in pooled_by_preferred_term(strata, outcome):
        if both:
            total += sum(both.values())
            largest += max(both.values())
    return min(overall, total / largest) if largest else 1.0


# ---------------------------------------------------------------------------
# Estimators
# ---------------------------------------------------------------------------


def mantel_haenszel(rows, inflate=1.0):
    """Mantel-Haenszel risk ratio with the Greenland-Robins variance, times `inflate`."""
    R = S = V = 0.0
    a_tot = c_tot = n1_tot = n0_tot = 0
    informative = 0
    for _, a, n1, c, n0 in rows:
        a_tot += a; c_tot += c; n1_tot += n1; n0_tot += n0
        if a + c == 0:
            continue
        informative += 1
        N = n1 + n0
        R += a * n0 / N
        S += c * n1 / N
        V += (n1 * n0 * (a + c) - a * c * N) / N ** 2
    result = {"glp1_affected": a_tot, "glp1_at_risk": n1_tot, "placebo_affected": c_tot, "placebo_at_risk": n0_tot,
              "strata_with_events": informative}
    if R == 0 or S == 0:
        result.update({"rr": None, "lo": None, "hi": None, "p": None})
        return result
    log_rr = math.log(R / S)
    se = math.sqrt(V / (R * S) * inflate)
    z = log_rr / se
    result.update({"rr": math.exp(log_rr), "lo": math.exp(log_rr - 1.959964 * se), "hi": math.exp(log_rr + 1.959964 * se),
                   "log_rr": log_rr, "se": se, "z": z, "p": 2 * stats.norm.sf(abs(z))})
    return result


def peto(rows):
    """Peto's one-step odds ratio."""
    OE = Vsum = 0.0
    for _, a, n1, c, n0 in rows:
        N, m = n1 + n0, a + c
        if m == 0 or m == N or N < 2:
            continue
        OE += a - n1 * m / N
        Vsum += n1 * n0 * m * (N - m) / (N ** 2 * (N - 1))
    if Vsum == 0:
        return {"or": None, "lo": None, "hi": None, "p": None}
    log_or, se = OE / Vsum, 1 / math.sqrt(Vsum)
    return {"or": math.exp(log_or), "lo": math.exp(log_or - 1.959964 * se), "hi": math.exp(log_or + 1.959964 * se),
            "p": 2 * stats.norm.sf(abs(log_or / se))}


def dersimonian_laird(rows):
    """Random-effects risk ratio (DerSimonian-Laird) on per-stratum log risk ratios, adding 0.5
    to every cell of a stratum with a zero cell, and leaving out strata with no events."""
    y, v = [], []
    for _, a, n1, c, n0 in rows:
        if a + c == 0:
            continue
        if a == 0 or c == 0 or a == n1 or c == n0:
            a, c, n1, n0 = a + 0.5, c + 0.5, n1 + 1, n0 + 1
        y.append(math.log((a / n1) / (c / n0)))
        v.append(1 / a - 1 / n1 + 1 / c - 1 / n0)
    if len(y) < 2:
        return {"rr": None, "lo": None, "hi": None, "p": None, "tau2": None, "i2": None}
    y, v = np.array(y), np.array(v)
    w = 1 / v
    mu_fe = np.sum(w * y) / np.sum(w)
    Q = float(np.sum(w * (y - mu_fe) ** 2))
    df = len(y) - 1
    C = np.sum(w) - np.sum(w ** 2) / np.sum(w)
    tau2 = max(0.0, (Q - df) / C)
    ws = 1 / (v + tau2)
    mu = float(np.sum(ws * y) / np.sum(ws))
    se = float(math.sqrt(1 / np.sum(ws)))
    return {"rr": math.exp(mu), "lo": math.exp(mu - 1.959964 * se), "hi": math.exp(mu + 1.959964 * se),
            "p": 2 * stats.norm.sf(abs(mu / se)), "tau2": tau2, "i2": max(0.0, (Q - df) / Q) if Q > 0 else 0.0,
            "q_heterogeneity_p": float(stats.chi2.sf(Q, df))}


def benjamini_hochberg(pvalues):
    """BH q-values for a list that may hold None (left as None)."""
    idx = [i for i, p in enumerate(pvalues) if p is not None]
    q = [None] * len(pvalues)
    m = len(idx)
    order = sorted(idx, key=lambda i: pvalues[i])
    running = 1.0
    for rank in range(m, 0, -1):
        i = order[rank - 1]
        running = min(running, pvalues[i] * m / rank)
        q[i] = running
    return q


# ---------------------------------------------------------------------------
# The analysis
# ---------------------------------------------------------------------------


def halves(strata):
    """Trials paired by size, one of each pair to the discovery half at random (seeded)."""
    sizes = {}
    for s in strata:
        sizes[s["nct"]] = sizes.get(s["nct"], 0) + s["glp1"].at_risk + s["placebo"].at_risk
    order = sorted(sizes, key=lambda n: (-sizes[n], n))
    rng = np.random.default_rng(P["split_seed"])
    half = {}
    for i in range(0, len(order), 2):
        pair = order[i:i + 2]
        if len(pair) == 1:
            half[pair[0]] = 1
            continue
        first = int(rng.integers(2))
        half[pair[first]] = 1
        half[pair[1 - first]] = 2
    return half


def prefilter():
    """The structured filter before screening: every study the searches found, kept if randomized,
    parallel or factorial, with at least two adverse-event groups. The screened set must be exactly
    what it keeps."""
    search = load("data/search.json")
    found = sorted({n for ids in search["results"].values() for n in ids})
    kept, excluded = [], {}
    for nct in found:
        with open(os.path.join(ROOT, "data", "ctgov", f"{nct}.json"), encoding="utf-8") as f:
            record = json.load(f)
        info = record["protocolSection"].get("designModule", {}).get("designInfo", {})
        groups = record.get("resultsSection", {}).get("adverseEventsModule", {}).get("eventGroups", [])
        if info.get("allocation") != "RANDOMIZED":
            reason = "not_randomized"
        elif info.get("interventionModel") not in ("PARALLEL", "FACTORIAL"):
            reason = "design_" + str(info.get("interventionModel")).lower()
        elif len(groups) < 2:
            reason = "fewer_than_two_adverse_event_groups"
        else:
            kept.append(nct)
            continue
        excluded[reason] = excluded.get(reason, 0) + 1
    assert kept == sorted(SCREENING), "the screened trials must be exactly those the structured filter keeps"
    return {"found": len(found), "excluded": excluded}


def overcount(strata):
    """How far summing term counts overstates participants with a serious event, over all strata
    and both arms: (sum over terms of participants affected) / (participants with any serious
    event), at least 1."""
    terms = sum(sum(s[arm].terms.values()) for s in strata for arm in ("glp1", "placebo"))
    people = sum(s[arm].any_serious for s in strata for arm in ("glp1", "placebo"))
    return max(1.0, terms / people)


def total_events(rows):
    return sum(a + c for _, a, _, c, _ in rows)


def estimate(strata, outcome, half_of, overall, with_sensitivity=True):
    rows = tables(strata, outcome)
    inflate = overlap_bound(strata, outcome, overall) if outcome not in ("death", "any_serious") else 1.0
    main = mantel_haenszel(rows)
    out = {"main": main, "events": total_events(rows)}
    if out["events"] < P["min_events"]:
        out["testable"] = False
        return out
    out["testable"] = True
    out["overlap_bound"] = inflate
    # When every participant affected is in one arm the risk ratio is undefined; Peto's odds
    # ratio, which needs no event in both arms, tests the outcome then.
    out["peto"] = peto(rows)
    out["p_primary"] = main["p"] if main["p"] is not None else out["peto"]["p"]
    out["higher_on_glp1"] = (main["rr"] > 1) if main["rr"] is not None else (out["peto"]["or"] > 1)
    for h in (1, 2):
        out[f"half{h}"] = mantel_haenszel([r for r in rows if half_of[r[0]["nct"]] == h])
    if not with_sensitivity:
        out["conservative"] = mantel_haenszel(rows, inflate)
        return out
    out["conservative"] = mantel_haenszel(rows, inflate)
    out["random_effects"] = dersimonian_laird(rows)
    if outcome not in ("death", "any_serious"):
        term = dominant_term(strata, outcome)
        out["dominant_term"] = {"term": term, **mantel_haenszel(tables(strata, outcome, term))}
    out["selective_only"] = mantel_haenszel([r for r in rows if r[0]["selective"]])
    out["type2_diabetes"] = mantel_haenszel([r for r in rows if r[0]["population"] == "type2_diabetes"])
    out["without_diabetes"] = mantel_haenszel([r for r in rows if r[0]["population"] == "overweight_obesity"])
    t2, ob = out["type2_diabetes"], out["without_diabetes"]
    if t2.get("se") and ob.get("se"):
        z = (t2["log_rr"] - ob["log_rr"]) / math.sqrt(t2["se"] ** 2 + ob["se"] ** 2)
        out["population_interaction_p"] = 2 * stats.norm.sf(abs(z))
    loo = []
    trials = sorted({r[0]["nct"] for r in rows if r[1] + r[3] > 0})
    for nct in trials:
        m = mantel_haenszel([r for r in rows if r[0]["nct"] != nct])
        if m["rr"] is not None:
            loo.append({"left_out": nct, "rr": m["rr"], "p": m["p"]})
    if loo:
        out["leave_one_out"] = {
            "rr_min": min(x["rr"] for x in loo), "rr_max": max(x["rr"] for x in loo),
            "p_max": max(x["p"] for x in loo), "p_max_left_out": max(loo, key=lambda x: x["p"])["left_out"],
        }
    return out


def power(rows, hr):
    """Power of a two-sided 5% test to detect a risk ratio `hr`, from the pooled count of
    participants with the outcome and the allocation, blind to which arm they were in."""
    events = total_events(rows)
    n1 = sum(r[2] for r in rows)
    n0 = sum(r[4] for r in rows)
    if events == 0 or n1 == 0 or n0 == 0:
        return 0.0
    f = n1 / (n1 + n0)
    se = math.sqrt(1 / (events * f) + 1 / (events * (1 - f)))
    return float(stats.norm.cdf(abs(math.log(hr)) / se - 1.959964))


def significant(p):
    return p is not None and p < P["alpha"]


def robust(e):
    """Significant (q < 0.05) and still so with the conservative variance, under random effects,
    and with every single trial left out."""
    if not e.get("testable"):
        return False
    return (significant(e["q"]) and significant(e["conservative"]["p"]) and significant(e["random_effects"]["p"])
            and "leave_one_out" in e and significant(e["leave_one_out"]["p_max"]))


def va_category(e, hr, lo, hi):
    """Where randomized evidence leaves an association the VA atlas reported."""
    m = e["main"]
    if not e.get("testable"):
        return "too_few_events"
    same = e["higher_on_glp1"] == (hr > 1)
    if significant(e["q"]):
        return "confirmed" if same else "opposite"
    if m["rr"] is None or not (m["lo"] <= hr <= m["hi"]):
        return "excludes_va_estimate"
    return "inconclusive"


def round_floats(obj, digits=6):
    if isinstance(obj, np.generic):
        obj = obj.item()
    if isinstance(obj, float):
        return float(f"{obj:.{digits}g}")
    if isinstance(obj, dict):
        return {k: round_floats(v, digits) for k, v in obj.items()}
    if isinstance(obj, list):
        return [round_floats(v, digits) for v in obj]
    return obj


def write(name, obj):
    with open(os.path.join(OUT, name), "w", encoding="utf-8") as f:
        json.dump(round_floats(obj), f, indent=1, sort_keys=True, ensure_ascii=False)
        f.write("\n")


def main():
    os.makedirs(OUT, exist_ok=True)
    strata = build_strata()
    half_of = halves(strata)
    inflate = overcount(strata)

    # R1: the trials.
    funnel = prefilter()
    eligible = sorted({s["nct"] for s in strata})
    by_molecule, by_population = {}, {}
    for s in strata:
        n = s["glp1"].at_risk + s["placebo"].at_risk
        for m in s["molecules"]:
            by_molecule[m] = by_molecule.get(m, 0) + n
        by_population[s["population"]] = by_population.get(s["population"], 0) + n
    reasons = {}
    for row in SCREENING.values():
        if not row["eligible"]:
            reasons[row["exclusion_reason"]] = reasons.get(row["exclusion_reason"], 0) + 1
    write("R1.json", {
        "found": funnel["found"],
        "structural_exclusions": funnel["excluded"],
        "screened": len(SCREENING),
        "eligible_trials": len(eligible),
        "comparisons": len(strata),
        "glp1_participants": sum(s["glp1"].at_risk for s in strata),
        "placebo_participants": sum(s["placebo"].at_risk for s in strata),
        "participants": sum(s["glp1"].at_risk + s["placebo"].at_risk for s in strata),
        "participants_by_molecule": by_molecule,
        "participants_by_population": by_population,
        "exclusions": reasons,
        "half1_trials": sum(1 for n in eligible if half_of[n] == 1),
        "half2_trials": sum(1 for n in eligible if half_of[n] == 2),
        "overcount_factor": inflate,
        "blind": BLIND,
    })

    # R2 and R3: the pre-registered outcomes, families A and B.
    for family, name in (("A", "R2.json"), ("B", "R3.json")):
        defs = [o for o in OUTCOMES if o["family"] == family]
        results = {}
        for o in defs:
            key = o.get("count", o["id"])
            e = estimate(strata, key, half_of, inflate)
            e["name"] = o["name"]
            if "va_hr" in o:
                e["va"] = {"hr": o["va_hr"], "lo": o["va_lo"], "hi": o["va_hi"]}
                e["power_for_va_hr"] = power(tables(strata, key), o["va_hr"])
            results[o["id"]] = e
        q = benjamini_hochberg([results[o["id"]]["p_primary"] if results[o["id"]]["testable"] else None for o in defs])
        for o, qv in zip(defs, q):
            e = results[o["id"]]
            e["q"] = qv
            e["robust"] = robust(e)
            if "va" in e:
                e["va_category"] = va_category(e, o["va_hr"], o["va_lo"], o["va_hi"])
                e["informative"] = e["power_for_va_hr"] >= P["informative_power"]
            if "expected" in o:
                e["expected"] = o["expected"]
                e["as_expected"] = bool(significant(e.get("q")) and e["higher_on_glp1"] == (o["expected"] == "higher"))
        write(name, results)
        if family == "A":
            a_results = results

    # R4: every organ system, discovery and replication halves.
    socs = sorted({soc for s in strata for arm in ("glp1", "placebo") for soc in s[arm].socs})
    scan = {}
    for soc in socs:
        e = estimate(strata, ("soc", soc), half_of, inflate, with_sensitivity=False)
        if e["testable"]:
            scan[soc] = e
    write("R4.json", scan_family(scan, inflate, strata, half_of, "soc"))

    # R5: every preferred term with enough participants affected, the same way.
    pts = {}
    for canonical in sorted(CANONICAL_TERMS):
        rows = tables(strata, ("pt", canonical))
        if total_events(rows) >= P["min_events_term_scan"]:
            pts[canonical] = estimate(strata, ("pt", canonical), half_of, 1.0, with_sensitivity=False)
    write("R5.json", scan_family(pts, 1.0, strata, half_of, "pt"))

    # R6: the VA atlas against randomized evidence.
    tally = {}
    informative = {}
    for oid, e in a_results.items():
        tally[e["va_category"]] = tally.get(e["va_category"], 0) + 1
        if e.get("informative"):
            informative[e["va_category"]] = informative.get(e["va_category"], 0) + 1
    write("R6.json", {
        "outcomes": len(a_results),
        "testable": sum(1 for e in a_results.values() if e["testable"]),
        "informative": sum(1 for e in a_results.values() if e.get("informative")),
        "categories": tally,
        "categories_informative": informative,
        "robust": sorted(oid for oid, e in a_results.items() if e["robust"]),
    })

    write_tables(strata, a_results)


def scan_family(results, inflate, strata, half_of, kind):
    """Discovery in half 1 (BH over the family, conservative variance), replication in half 2
    (p < 0.05 in the same direction), and the whole for reference."""
    names = sorted(results)
    bound = {n: (overlap_bound(strata, (kind, n), inflate) if kind == "soc" else 1.0) for n in names}
    disc = [mantel_haenszel([r for r in tables(strata, (kind, n)) if half_of[r[0]["nct"]] == 1], bound[n]) for n in names]
    rep = [mantel_haenszel([r for r in tables(strata, (kind, n)) if half_of[r[0]["nct"]] == 2], bound[n]) for n in names]
    q = benjamini_hochberg([d["p"] for d in disc])
    out = {"tested": len(names), "items": {}}
    discovered, replicated = [], []
    for n, d, r, qv in zip(names, disc, rep, q):
        item = {"all": results[n]["main"], "discovery": d, "replication": r, "discovery_q": qv, "overlap_bound": bound[n]}
        item["discovered"] = qv is not None and qv < P["alpha"]
        item["replicated"] = bool(item["discovered"] and r["p"] is not None and r["p"] < P["alpha"]
                                  and (r["rr"] > 1) == (d["rr"] > 1))
        if item["discovered"]:
            discovered.append(n)
        if item["replicated"]:
            replicated.append(n)
        out["items"][n] = item
    out["discovered"] = discovered
    out["replicated"] = replicated
    return out


def write_tables(strata, a_results):
    """Per-stratum counts for every family A and B outcome, for readers to check by hand."""
    with open(os.path.join(OUT, "strata.csv"), "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["stratum", "nct", "molecules", "population", "glp1_at_risk", "placebo_at_risk",
                    "glp1_any_serious", "placebo_any_serious"])
        for s in strata:
            w.writerow([s["id"], s["nct"], "+".join(s["molecules"]), s["population"], s["glp1"].at_risk,
                        s["placebo"].at_risk, s["glp1"].any_serious, s["placebo"].any_serious])
    with open(os.path.join(OUT, "outcome_counts.csv"), "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["outcome", "stratum", "glp1_affected", "glp1_at_risk", "placebo_affected", "placebo_at_risk"])
        for o in OUTCOMES:
            for s, a, n1, c, n0 in tables(strata, o.get("count", o["id"])):
                if a + c:
                    w.writerow([o["id"], s["id"], a, n1, c, n0])


if __name__ == "__main__":
    main()