"""Serious suicidal events in randomized placebo-controlled trials of antidepressants.

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 and
analysis definitions (data/outcomes.json, data/analyses.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 re
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")
Z95 = 1.959964


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 = {o["id"]: o for o in load("data/outcomes.json")}
ANALYSES = load("data/analyses.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")}


def record(nct):
    with open(os.path.join(ROOT, "data", "ctgov", f"{nct}.json"), encoding="utf-8") as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# Trial facts from the record's structured fields
# ---------------------------------------------------------------------------


def years(text):
    """An eligibility age such as "18 Years" or "6 Months", in years; None if not given."""
    if not text:
        return None
    m = re.match(r"\s*(\d+(?:\.\d+)?)\s*(\w+)", text)
    if not m:
        return None
    value, unit = float(m.group(1)), m.group(2).lower()
    per_year = {"year": 1, "years": 1, "month": 12, "months": 12, "week": 52.1775, "weeks": 52.1775, "day": 365.25, "days": 365.25}
    return value / per_year[unit] if unit in per_year else None


def age_group(rec):
    """From the eligible ages: pediatric (all under 18), older (all 60 or over), adult (all 18 or
    over), or mixed (spanning 18, or not stated)."""
    el = rec["protocolSection"].get("eligibilityModule", {})
    lo, hi = years(el.get("minimumAge")), years(el.get("maximumAge"))
    if hi is not None and hi < 18:
        return "pediatric"
    if lo is not None and lo >= 60:
        return "older"
    if lo is not None and lo >= 18:
        return "adult"
    return "mixed"


def mean_age(rec):
    """The trial's mean (or, failing that, median) age at baseline in years, from the posted
    "Age, Continuous" measure: the total group's value, or else the participant-weighted mean of
    the groups' values. None if the record has no such measure."""
    bl = rec.get("resultsSection", {}).get("baselineCharacteristicsModule", {})
    groups = {g["id"]: norm(g.get("title")) for g in bl.get("groups", [])}
    counts = {}
    for d in bl.get("denoms", []):
        for c in d.get("counts", []):
            try:
                counts[c["groupId"]] = float(c["value"])
            except (TypeError, ValueError):
                pass
    for kind in ("MEAN", "MEDIAN"):
        for m in bl.get("measures", []):
            if re.sub(r"[^a-z]", "", m.get("title", "").lower()) != "agecontinuous" or m.get("paramType") != kind:
                continue
            if norm(m.get("unitOfMeasure")) not in ("years", "year"):
                continue
            values = {}
            for cl in m.get("classes", []):
                for cat in cl.get("categories", []):
                    for x in cat.get("measurements", []):
                        try:
                            values[x["groupId"]] = float(x["value"])
                        except (TypeError, ValueError, KeyError):
                            pass
            total = [g for g, t in groups.items() if t == "total" and g in values]
            if total:
                return values[total[0]]
            parts = [(values[g], counts.get(g, 0)) for g in values if groups.get(g) != "total"]
            weight = sum(n for _, n in parts)
            if parts and weight > 0:
                return sum(v * n for v, n in parts) / weight
    return None


def completion(rec):
    """The primary completion date as YYYY-MM, or the study completion date if that is missing."""
    sm = rec["protocolSection"].get("statusModule", {})
    for key in ("primaryCompletionDateStruct", "completionDateStruct"):
        date = (sm.get(key) or {}).get("date")
        if date:
            return date[:7]
    return None


# ---------------------------------------------------------------------------
# Comparisons: each a stratum with an antidepressant arm and a placebo arm, pooled over groups
# ---------------------------------------------------------------------------


class Arm:
    def __init__(self):
        self.at_risk = 0  # participants at risk of a serious event
        self.other_at_risk = 0  # participants at risk of a non-serious event
        self.any_serious = 0
        self.deaths = 0
        self.deaths_at_risk = 0
        self.deaths_known = True
        self.terms = {"serious": {}, "other": {}}  # normalized term -> participants affected
        self.listed = {"serious": set(), "other": set()}  # terms the trial's tables list at all


def group_at_risk(group, rows, key):
    if group.get(key):
        return group[key]
    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("deathsNumAtRisk") or 0


def build_arm(rec, group_ids):
    ae = rec["resultsSection"]["adverseEventsModule"]
    groups = {g["id"]: g for g in ae["eventGroups"]}
    tables = {"serious": ae.get("seriousEvents", []), "other": ae.get("otherEvents", [])}
    arm = Arm()
    for gid in group_ids:
        g = groups[gid]
        arm.at_risk += group_at_risk(g, tables["serious"], "seriousNumAtRisk")
        arm.other_at_risk += group_at_risk(g, tables["other"], "otherNumAtRisk")
        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 table, rows in tables.items():
        for e in rows:
            term = norm(e.get("term"))
            arm.listed[table].add(term)
            n = sum(s.get("numAffected") or 0 for s in e.get("stats", []) if s.get("groupId") in wanted)
            if n:
                arm.terms[table][term] = arm.terms[table].get(term, 0) + n
    return arm


def build_strata():
    strata = []
    for nct in sorted(SCREENING):
        row = SCREENING[nct]
        if not row["eligible"]:
            continue
        rec = record(nct)
        facts = {"age_group": age_group(rec), "mean_age": mean_age(rec), "completion": completion(rec)}
        for k, comp in enumerate(row["comparisons"]):
            drug = build_arm(rec, comp["antidepressant_groups"])
            placebo = build_arm(rec, comp["placebo_groups"])
            if drug.at_risk == 0 or placebo.at_risk == 0:
                continue
            strata.append({
                "id": f"{nct}#{k + 1}",
                "nct": nct,
                "drug": drug,
                "placebo": placebo,
                "molecules": sorted(comp["molecules"]),
                "classes": sorted({MOLECULES[m] for m in comp["molecules"]}),
                "combination": bool(comp.get("combination")),
                "design": row["design"],
                "indication": row["indication"],
                **facts,
            })
    if BLIND:
        rng = np.random.default_rng(P["blind_seed"])
        for s in strata:
            if rng.random() < 0.5:
                s["drug"], s["placebo"] = s["placebo"], s["drug"]
    return strata


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

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


def outcome_terms(outcome):
    """The normalized terms an outcome counts: every term coded to any of its codes."""
    return set().union(*(TERMS_OF.get(c, set()) for c in OUTCOMES[outcome]["codes"]))


def by_preferred_term(arm, outcome):
    """An outcome's participants affected in an arm, by preferred term: {canonical: affected}."""
    table = OUTCOMES[outcome]["table"]
    out = {}
    for term in outcome_terms(outcome):
        n = arm.terms[table].get(term, 0)
        if n:
            canonical = CODING[term]["canonical"]
            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 when the trial's
    tables can't tell. Serious outcomes are counted in every comparison (a serious table lists
    every serious event). A non-serious outcome is counted only where the trial's non-serious
    table lists one of its terms, since that table omits terms below its frequency threshold."""
    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
    o = OUTCOMES[outcome]
    if o["table"] == "other":
        if not (outcome_terms(outcome) & arm.listed["other"]):
            return None
        at_risk = arm.other_at_risk
    else:
        at_risk = 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, at_risk), at_risk


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


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


# ---------------------------------------------------------------------------
# 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 = {"drug_affected": a_tot, "drug_at_risk": n1_tot, "placebo_affected": c_tot, "placebo_at_risk": n0_tot,
              "comparisons": len(rows), "comparisons_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 - Z95 * se), "hi": math.exp(log_rr + Z95 * se),
                   "log_rr": log_rr, "se": se, "z": z, "p": 2 * stats.norm.sf(abs(z))})
    return result


def peto_terms(a, n1, c, n0):
    """Peto's O - E and V for one stratum, or None if it carries no information."""
    N, m = n1 + n0, a + c
    if m == 0 or m == N or N < 2:
        return None
    return a - n1 * m / N, n1 * n0 * m * (N - m) / (N ** 2 * (N - 1))


def peto(rows):
    """Peto's one-step odds ratio."""
    OE = Vsum = 0.0
    for _, a, n1, c, n0 in rows:
        t = peto_terms(a, n1, c, n0)
        if t:
            OE += t[0]
            Vsum += t[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 - Z95 * se), "hi": math.exp(log_or + Z95 * 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 - Z95 * se), "hi": math.exp(mu + Z95 * 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


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


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


def halves(strata):
    """Trials paired by size, one of each pair to the first half at random (seeded)."""
    sizes = {}
    for s in strata:
        sizes[s["nct"]] = sizes.get(s["nct"], 0) + s["drug"].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 overcount(strata):
    """How far summing serious term counts overstates participants with a serious event, over all
    strata and both arms, at least 1."""
    terms = sum(sum(s[arm].terms["serious"].values()) for s in strata for arm in ("drug", "placebo"))
    people = sum(s[arm].any_serious for s in strata for arm in ("drug", "placebo"))
    return max(1.0, terms / people) if people else 1.0


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, capped at `overall`."""
    total = largest = 0
    for s in strata:
        both = by_preferred_term(s["drug"], outcome)
        for term, n in by_preferred_term(s["placebo"], outcome).items():
            both[term] = both.get(term, 0) + n
        if both:
            total += sum(both.values())
            largest += max(both.values())
    return min(overall, total / largest) if largest else 1.0


def dominant_term(strata, outcome):
    totals = {}
    for s in strata:
        for arm in ("drug", "placebo"):
            for term, n in by_preferred_term(s[arm], outcome).items():
                totals[term] = totals.get(term, 0) + n
    return min(totals, key=lambda t: (-totals[t], t)) if totals else None


def select(strata, spec):
    """The comparisons an analysis covers, by its registered filters."""
    out = []
    for s in strata:
        if "design" in spec and s["design"] not in spec["design"]:
            continue
        if "age_group" in spec and s["age_group"] not in spec["age_group"]:
            continue
        if "indication" in spec and s["indication"] not in spec["indication"]:
            continue
        if not spec.get("combinations", False) and s["combination"]:
            continue
        if "classes" in spec and not set(s["classes"]) <= set(spec["classes"]):
            continue
        if "completed_after" in spec and not (s["completion"] and s["completion"] > spec["completed_after"]):
            continue
        out.append(s)
    return out


def estimate(rows_strata, outcome, half_of, overall, full=True):
    rows = tables(rows_strata, outcome)
    main = mantel_haenszel(rows)
    out = {"main": main, "events": total_events(rows), "peto": peto(rows)}
    out["testable"] = out["events"] >= P["min_events"]
    out["p_primary"] = main["p"] if main["p"] is not None else out["peto"]["p"]
    if main["rr"] is not None:
        out["higher_on_drug"] = main["rr"] > 1
    elif out["peto"]["or"] is not None:
        out["higher_on_drug"] = out["peto"]["or"] > 1
    else:
        out["higher_on_drug"] = None
    if not full or not out["testable"]:
        return out
    inflate = overlap_bound(rows_strata, outcome, overall) if outcome not in ("death", "any_serious") else 1.0
    out["overlap_bound"] = inflate
    out["conservative"] = mantel_haenszel(rows, inflate)
    out["random_effects"] = dersimonian_laird(rows)
    for h in (1, 2):
        out[f"half{h}"] = mantel_haenszel([r for r in rows if half_of[r[0]["nct"]] == h])
    if outcome not in ("death", "any_serious"):
        term = dominant_term(rows_strata, outcome)
        out["dominant_term"] = {"term": term, **mantel_haenszel(tables(rows_strata, outcome, term))}
    loo = []
    for nct in sorted({r[0]["nct"] for r in rows if r[1] + r[3] > 0}):
        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"], "lo": m["lo"], "hi": m["hi"]})
    if loo:
        out["leave_one_out"] = {
            "rr_min": min(x["rr"] for x in loo), "rr_max": max(x["rr"] for x in loo),
            "hi_max": max(x["hi"] for x in loo), "lo_min": min(x["lo"] for x in loo),
            "p_min": min(x["p"] for x in loo), "p_max": max(x["p"] for x in loo),
        }
    return out


def power(rows, rr):
    """Power of a two-sided 5% test to detect a risk ratio `rr`, 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(rr)) / se - Z95))


def against_priors(m, compare):
    """Where an estimate's interval sits against the earlier estimates registered for it."""
    if m["rr"] is None:
        return None
    out = {"increase": m["lo"] > 1, "decrease": m["hi"] < 1}
    for name in compare:
        value = P["priors"][name]
        out[f"excludes_{name}"] = not (m["lo"] <= value <= m["hi"])
    return out


def holm(pvalues):
    """Holm-adjusted p-values for a list that may hold None (left as None)."""
    idx = sorted((i for i, p in enumerate(pvalues) if p is not None), key=lambda i: pvalues[i])
    m, out, running = len(idx), [None] * len(pvalues), 0.0
    for rank, i in enumerate(idx):
        running = max(running, min(1.0, (m - rank) * pvalues[i]))
        out[i] = running
    return out


def age_regression(rows):
    """Peto meta-regression of the log odds ratio on the trial's mean age: a weighted regression
    of (O - E)/V on mean age with weights V, over comparisons with events and a known mean age.
    Returns the slope per year of age and the implied percentage change in the odds ratio."""
    pts = []
    for s, a, n1, c, n0 in rows:
        t = peto_terms(a, n1, c, n0)
        if t and s["mean_age"] is not None:
            pts.append((s["mean_age"], t[0] / t[1], t[1]))
    if len(pts) < 3:
        return {"comparisons": len(pts), "slope": None}
    x = np.array([p[0] for p in pts]); y = np.array([p[1] for p in pts]); w = np.array([p[2] for p in pts])
    xbar = np.sum(w * x) / np.sum(w)
    sxx = float(np.sum(w * (x - xbar) ** 2))
    if sxx == 0:
        return {"comparisons": len(pts), "slope": None}
    slope = float(np.sum(w * (x - xbar) * y) / sxx)
    se = 1 / math.sqrt(sxx)
    return {"comparisons": len(pts), "slope": slope, "se": se, "p": 2 * stats.norm.sf(abs(slope / se)),
            "pct_per_year": 100 * (math.exp(slope) - 1), "pct_lo": 100 * (math.exp(slope - Z95 * se) - 1),
            "pct_hi": 100 * (math.exp(slope + Z95 * se) - 1), "age_min": float(x.min()), "age_max": float(x.max()),
            "age_weighted_mean": float(xbar)}


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:
        rec = record(nct)
        info = rec["protocolSection"].get("designModule", {}).get("designInfo", {})
        groups = rec.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 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 describe(strata):
    n = lambda ss: sum(s["drug"].at_risk + s["placebo"].at_risk for s in ss)
    by = lambda key: {v: n([s for s in strata if s[key] == v]) for v in sorted({s[key] for s in strata})}
    by_class = {}
    for s in strata:
        for c in s["classes"]:
            by_class[c] = by_class.get(c, 0) + s["drug"].at_risk + s["placebo"].at_risk
    return {
        "trials": len({s["nct"] for s in strata}), "comparisons": len(strata), "participants": n(strata),
        "drug_participants": sum(s["drug"].at_risk for s in strata), "placebo_participants": sum(s["placebo"].at_risk for s in strata),
        "by_design": by("design"), "by_indication": by("indication"), "by_age_group": by("age_group"), "by_class": by_class,
        "with_mean_age": sum(1 for s in strata if s["mean_age"] is not None),
    }


def main():
    os.makedirs(OUT, exist_ok=True)
    funnel = prefilter()
    strata = build_strata()
    half_of = halves(strata)
    overall = overcount(strata)
    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({s["nct"] for s in strata}), "exclusions": reasons, "overcount_factor": overall,
                      "all": describe(strata), "blind": BLIND,
                      "analyses": {a["id"]: describe(select(strata, a["select"])) for a in ANALYSES}})

    # R2: the co-primary analyses, in full, Holm-adjusted across the two.
    primaries = [a for a in ANALYSES if a["role"] == "primary"]
    results = {}
    for a in primaries:
        sel = select(strata, a["select"])
        e = estimate(sel, a["outcome"], half_of, overall)
        rows = tables(sel, a["outcome"])
        e["against_priors"] = against_priors(e["main"], a.get("compare", []))
        e["power"] = {name: power(rows, P["priors"][name]) for name in a.get("compare", [])}
        e["robust"] = bool(significant(e["p_primary"]) 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"]))
        e["outcome"], e["name"] = a["outcome"], a["name"]
        results[a["id"]] = e
    for a, ph in zip(primaries, holm([results[a["id"]]["p_primary"] for a in primaries])):
        results[a["id"]]["p_holm"] = ph
    write("R2.json", results)

    # R3: the secondary analyses, Benjamini-Hochberg over the family.
    secondary = [a for a in ANALYSES if a["role"] == "secondary"]
    results = {}
    for a in secondary:
        sa = select(strata, a["select"])
        r = estimate(sa, a["outcome"], half_of, overall)
        r["name"], r["outcome"] = a["name"], a["outcome"]
        r["against_priors"] = against_priors(r["main"], a.get("compare", []))
        r["power"] = {name: power(tables(sa, a["outcome"]), P["priors"][name]) for name in a.get("compare", [])}
        results[a["id"]] = r
    q = benjamini_hochberg([results[a["id"]]["p_primary"] if results[a["id"]]["testable"] else None for a in secondary])
    for a, qv in zip(secondary, q):
        results[a["id"]]["q"] = qv
    write("R3.json", results)

    # R4: age, as the trials' mean ages vary.
    regressions = {}
    for a in [a for a in ANALYSES if a["role"] == "age"]:
        regressions[a["id"]] = {"name": a["name"], "outcome": a["outcome"],
                                **age_regression(tables(select(strata, a["select"]), a["outcome"]))}
    write("R4.json", regressions)

    # R5: checks of the data.
    checks = {}
    for a in [a for a in ANALYSES if a["role"] == "check"]:
        r = estimate(select(strata, a["select"]), a["outcome"], half_of, overall, full=False)
        r["name"], r["outcome"], r["expected"] = a["name"], a["outcome"], a.get("expected")
        if a.get("expected"):
            r["as_expected"] = bool(significant(r["p_primary"]) and r["higher_on_drug"] == (a["expected"] == "higher"))
        checks[a["id"]] = r
    write("R5.json", checks)

    with open(os.path.join(OUT, "comparisons.csv"), "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["comparison", "nct", "molecules", "design", "indication", "age_group", "mean_age", "completion", "half",
                    "drug_at_risk", "placebo_at_risk", "drug_suicidal_behaviour", "placebo_suicidal_behaviour",
                    "drug_suicidal_ideation", "placebo_suicidal_ideation"])
        for s in strata:
            b = [count(s[arm], "S1") for arm in ("drug", "placebo")]
            i = [count(s[arm], "S2") for arm in ("drug", "placebo")]
            w.writerow([s["id"], s["nct"], "+".join(s["molecules"]), s["design"], s["indication"], s["age_group"],
                        "" if s["mean_age"] is None else round(s["mean_age"], 1), s["completion"], half_of[s["nct"]],
                        s["drug"].at_risk, s["placebo"].at_risk, b[0][0], b[1][0], i[0][0], i[1][0]])


if __name__ == "__main__":
    main()