"""Analyses added after the first version's reviews, beside the registered ones (deviations.json).

The registered analysis (code/analyze.py) is unchanged and still writes results/R1.json to R6.json.
This script reads the same strata through it and writes results/V.json:

- what each outcome's count is: participants (one preferred term, deaths, any serious event) or
  reports (several terms summed, a participant counted once for each term they had);
- a variance clustered by trial, which with many independent trials holds however terms overlap
  within a participant, and estimates that use the wider of it and the registered Greenland-Robins
  variance, with a revised robustness rule on them; the random-effects leg of that rule is still
  the registered one, on binomial variances;
- the registered scans of organ systems and preferred terms again, with the wider variance;
- how far overlap between terms could move the estimates: bounds from the tables alone, each
  comparison arm's count divided by its own overlap across all serious events, the same with the
  overlap pooled over all arms, and the uniform difference in overlap between arms that would bring
  the unadjusted interval to one;
- the preferred terms within bacterial infections and pneumonia, each counted in participants
  except where a group lists one term in several rows, and pneumonia without coronavirus pneumonia;
- the atlas against the trials as two estimates, each with its uncertainty, corrected over the
  outcomes tested;
- every estimate again by how long the trials recorded serious events (data/ae_window.json).

Usage: python3 code/revised.py <bundle root>
"""

import json
import math
import os

from scipy import stats

import analyze as A

ROOT = A.ROOT
REVISED = A.load("data/revised.json")
WINDOW = {row["nct"]: row["window"] for row in A.load("data/ae_window.json")}
WINDOWS = ("on_treatment", "in_trial", "calendar", "unstated")
Z = 1.959964


# ---------------------------------------------------------------------------
# Counting
# ---------------------------------------------------------------------------


def terms_counted(strata, outcome):
    """The preferred terms an outcome's count adds up, over all strata and both arms."""
    if outcome in ("death", "any_serious"):
        return []
    seen = set()
    for s in strata:
        for arm in ("glp1", "placebo"):
            seen.update(A.by_preferred_term(s[arm], outcome))
    return sorted(seen)


def tables_without(strata, outcome, drop):
    """A.tables for an outcome with some of its preferred terms left out."""
    out = []
    for s in strata:
        cells = []
        for arm in ("glp1", "placebo"):
            counts = A.by_preferred_term(s[arm], outcome)
            cells.append((min(sum(n for t, n in counts.items() if t not in drop), s[arm].at_risk), s[arm].at_risk))
        out.append((s, cells[0][0], cells[0][1], cells[1][0], cells[1][1]))
    return out


def in_outcome(outcome, term, soc):
    """Whether a reported row, by its normalized term and organ system, counts toward an outcome."""
    if outcome[0] == "soc":
        return soc == outcome[1]
    if outcome[0] == "pt":
        return term in A.CANONICAL_TERMS[outcome[1]]
    return term in A.TERMS_OF.get(outcome, ())


class Groups:
    """The adverse-event groups behind each comparison arm, row by row as the records report them."""

    def __init__(self, strata):
        self.records = {}
        self.groups = []  # (arm, record, group id)
        for s in strata:
            nct, k = s["id"].split("#")
            if nct not in self.records:
                with open(os.path.join(ROOT, "data", "ctgov", f"{nct}.json"), encoding="utf-8") as f:
                    self.records[nct] = json.load(f)
            comp = A.SCREENING[nct]["comparisons"][int(k) - 1]
            for arm, key in (("glp1", "glp1_groups"), ("placebo", "placebo_groups")):
                for gid in comp[key]:
                    self.groups.append((arm, self.records[nct], gid))

    @staticmethod
    def rows(record, gid, outcome):
        """Each reported row's participants affected in a group, for the rows that count toward an outcome."""
        out = []
        for e in record["resultsSection"]["adverseEventsModule"].get("seriousEvents", []):
            if in_outcome(outcome, A.norm(e.get("term")), A.norm(e.get("organSystem"))):
                n = sum(st.get("numAffected") or 0 for st in e.get("stats", []) if st.get("groupId") == gid)
                if n:
                    out.append(n)
        return out

    def bounds(self, outcome):
        """For each arm, the fewest and most participants the tables allow to have the outcome, summed
        over the adverse-event groups: at least its largest reported row, and at least its count less
        the group's surplus of term counts over participants with any serious event; at most its count,
        and at most the participants with any serious event."""
        out = {"glp1": [0, 0], "placebo": [0, 0]}
        for arm, record, gid in self.groups:
            counts = self.rows(record, gid, outcome)
            if not counts:
                continue
            g = A.build_arm(record, [gid])
            total = sum(counts)
            surplus = max(0, sum(g.terms.values()) - g.any_serious)
            out[arm][0] += max(max(counts), total - surplus)
            out[arm][1] += min(total, g.any_serious) if g.any_serious else total
        return {arm: {"fewest": lo, "most": hi} for arm, (lo, hi) in out.items()}

    def merged_rows(self, term):
        """Groups that report a preferred term in more than one row (the same term under two organ
        systems, or terms the coding gave one preferred term), and the participants in all but the
        largest row of each: how far the term's count could exceed its participants."""
        groups = extra = 0
        for _, record, gid in self.groups:
            counts = self.rows(record, gid, ("pt", term))
            if len(counts) > 1:
                groups += 1
                extra += sum(counts) - max(counts)
        return {"groups": groups, "participants": extra}


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


def clustered_variance(rows):
    """Variance of the log Mantel-Haenszel risk ratio from each trial's contribution (a sandwich
    over trials, times G / (G - 1) for G trials with events). It treats trials as independent and
    leaves how a trial's count arises free, so with many trials it holds when a participant is
    counted once for each term; with few trials, or a few carrying most of the weight, it can be
    too small."""
    R = S = 0.0
    per = {}
    for s, a, n1, c, n0 in rows:
        N = n1 + n0
        r, q = a * n0 / N, c * n1 / N
        R += r
        S += q
        cell = per.setdefault(s["nct"], [0.0, 0.0])
        cell[0] += r
        cell[1] += q
    if R == 0 or S == 0:
        return None
    G = sum(1 for r, q in per.values() if r + q > 0)
    if G < 2:
        return None
    return sum((r / R - q / S) ** 2 for r, q in per.values()) * G / (G - 1)


def wider(rows):
    """The Mantel-Haenszel risk ratio with the wider of the Greenland-Robins and the clustered
    variances, and each of them."""
    m = A.mantel_haenszel(rows)
    out = {k: m[k] for k in ("glp1_affected", "glp1_at_risk", "placebo_affected", "placebo_at_risk", "strata_with_events")}
    out["trials"] = len({r[0]["nct"] for r in rows})
    out["trials_with_events"] = len({r[0]["nct"] for r in rows if r[1] + r[3] > 0})
    if m["rr"] is None:
        out.update({"rr": None, "lo": None, "hi": None, "p": None, "se": None})
        return out
    v_gr = m["se"] ** 2
    v_cl = clustered_variance(rows)
    se = math.sqrt(max(v_gr, v_cl if v_cl is not None else 0.0))
    out.update({
        "rr": m["rr"], "log_rr": m["log_rr"], "se": se,
        "lo": math.exp(m["log_rr"] - Z * se), "hi": math.exp(m["log_rr"] + Z * se),
        "p": 2 * stats.norm.sf(abs(m["log_rr"] / se)),
        "se_greenland_robins": m["se"], "se_clustered": None if v_cl is None else math.sqrt(v_cl),
        "clustered_wider": v_cl is not None and v_cl > v_gr,
    })
    return out


def difference(x, y):
    """Two-sided p that two independent log risk ratios differ, and their ratio with its interval."""
    if not x.get("se") or not y.get("se"):
        return None
    d, se = x["log_rr"] - y["log_rr"], math.hypot(x["se"], y["se"])
    return {"p": 2 * stats.norm.sf(abs(d) / se), "ratio": math.exp(d), "lo": math.exp(d - Z * se), "hi": math.exp(d + Z * se)}


def leave_one_out(rows):
    out = []
    for nct in sorted({r[0]["nct"] for r in rows if r[1] + r[3] > 0}):
        m = wider([r for r in rows if r[0]["nct"] != nct])
        if m["rr"] is not None:
            out.append({"left_out": nct, "rr": m["rr"], "p": m["p"]})
    if not out:
        return None
    worst = max(out, key=lambda x: x["p"])
    return {"rr_min": min(x["rr"] for x in out), "rr_max": max(x["rr"] for x in out), "p_max": worst["p"],
            "p_max_left_out": worst["left_out"]}


# ---------------------------------------------------------------------------
# Overlap between terms
# ---------------------------------------------------------------------------


def terms_per_person(arm):
    """How many serious terms an arm reports for each participant with a serious event, at least 1."""
    if not arm.any_serious:
        return 1.0
    return max(1.0, sum(arm.terms.values()) / arm.any_serious)


def overlap_adjusted(rows):
    """Each comparison arm's count divided by that arm's own terms per participant over all serious
    events: what the count would be if the outcome's terms overlapped within a participant as all
    serious terms do in that arm of that comparison."""
    return [(s, a / terms_per_person(s["glp1"]), n1, c / terms_per_person(s["placebo"]), n0) for s, a, n1, c, n0 in rows]


def overlap_pooled(rows, factors):
    """The same with one factor per arm pooled over every comparison, which mixes the differences
    between trials into the difference between arms."""
    return [(s, a / factors["glp1"], n1, c / factors["placebo"], n0) for s, a, n1, c, n0 in rows]


def tipping(m):
    """The ratio of one arm's terms per participant to the other's, the same in every comparison, at
    which the unadjusted interval would reach one, its width held fixed: placebo's over the drug's
    for an outcome lower on the drug, the drug's over placebo's for one higher."""
    if m["rr"] is None or m["lo"] <= 1 <= m["hi"]:
        return None
    return 1 / m["hi"] if m["hi"] < 1 else m["lo"]


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


def by_window(rows):
    groups = {w: [r for r in rows if WINDOW[r[0]["nct"]] == w] for w in WINDOWS}
    groups["not_on_treatment"] = [r for r in rows if WINDOW[r[0]["nct"]] != "on_treatment"]
    out = {k: wider(v) for k, v in groups.items()}
    d = difference(out["on_treatment"], out["not_on_treatment"])
    out["interaction_p"] = None if d is None else d["p"]
    return out


def scan(strata, half_of, kind, names):
    """A registered scan again with the wider variance: each item in the first half, Benjamini-Hochberg
    over the family, and found again when q < 0.05 there and p < 0.05 in the second half in the same
    direction."""
    items = {}
    for n in names:
        rows = A.tables(strata, (kind, n))
        items[n] = {"all": wider(rows),
                    "discovery": wider([r for r in rows if half_of[r[0]["nct"]] == 1]),
                    "replication": wider([r for r in rows if half_of[r[0]["nct"]] == 2])}
    q = A.benjamini_hochberg([items[n]["discovery"]["p"] for n in names])
    for n, qv in zip(names, q):
        it = items[n]
        d, r = it["discovery"], it["replication"]
        it["discovery_q"] = qv
        it["discovered"] = A.significant(qv)
        it["replicated"] = bool(it["discovered"] and A.significant(r["p"]) and (r["rr"] > 1) == (d["rr"] > 1))
    found = [n for n in names if items[n]["replicated"]]
    return {"tested": len(names), "items": items,
            "discovered": [n for n in names if items[n]["discovered"]],
            "replicated_lower": [n for n in found if items[n]["all"]["rr"] < 1],
            "replicated_higher": [n for n in found if items[n]["all"]["rr"] > 1]}


def main():
    strata = A.build_strata()
    groups = Groups(strata)
    registered = {**A.load("results/R2.json"), **A.load("results/R3.json")}
    arms = {arm: {"terms": sum(sum(s[arm].terms.values()) for s in strata),
                  "participants_with_any": sum(s[arm].any_serious for s in strata)} for arm in ("glp1", "placebo")}
    for a in arms.values():
        a["terms_per_participant"] = a["terms"] / a["participants_with_any"]
    pooled = {arm: a["terms_per_participant"] for arm, a in arms.items()}

    families = {}
    outcomes = {}
    for o in A.OUTCOMES:
        key = o.get("count", o["id"])
        rows = A.tables(strata, key)
        e = registered[o["id"]]
        terms = terms_counted(strata, key)
        out = {"name": o["name"], "family": o["family"], "testable": e["testable"], "preferred_terms": len(terms),
               "counts": "participants" if key in ("death", "any_serious") or len(terms) <= 1 else "reports",
               "main": wider(rows)}
        if e["testable"]:
            out["leave_one_out"] = leave_one_out(rows)
            out["window"] = by_window(rows)
            if out["counts"] == "reports":
                adjusted = wider(overlap_adjusted(rows))
                out["overlap"] = {"adjusted": adjusted, "pooled": wider(overlap_pooled(rows, pooled)),
                                  "within_trial_shift": None if adjusted["rr"] is None else adjusted["rr"] / out["main"]["rr"],
                                  "tipping_ratio": tipping(out["main"]), "bounds": groups.bounds(key)}
            if "va_hr" in o:
                atlas = {"log_rr": math.log(o["va_hr"]), "se": (math.log(o["va_hi"]) - math.log(o["va_lo"])) / (2 * Z)}
                d = difference(out["main"], atlas)
                out["atlas"] = {"hr": o["va_hr"], "lo": o["va_lo"], "hi": o["va_hi"],
                                "difference_p": None if d is None else d["p"],
                                "ratio": None if d is None else d["ratio"],
                                "ratio_lo": None if d is None else d["lo"], "ratio_hi": None if d is None else d["hi"],
                                "comparable": o["id"] not in REVISED["not_comparable"],
                                "why_not_comparable": REVISED["not_comparable"].get(o["id"])}
        outcomes[o["id"]] = out
        families.setdefault(o["family"], []).append(o["id"])

    # q over each family from the wider variance, and the revised robustness rule.
    for fam, ids in families.items():
        q = A.benjamini_hochberg([outcomes[i]["main"]["p"] if outcomes[i]["testable"] else None for i in ids])
        for i, qv in zip(ids, q):
            out = outcomes[i]
            out["q"] = qv
            if not out["testable"]:
                continue
            re_p = registered[i]["random_effects"]["p"]
            loo = out["leave_one_out"]
            out["robust"] = bool(A.significant(qv) and A.significant(re_p) and loo and A.significant(loo["p_max"]))

    # The atlas against the trials: the difference between the two estimates, corrected over
    # every family A outcome tested, and where each outcome then stands.
    tested = [i for i in families["A"] if outcomes[i]["testable"] and outcomes[i]["atlas"]["difference_p"] is not None]
    for i, qv in zip(tested, A.benjamini_hochberg([outcomes[i]["atlas"]["difference_p"] for i in tested])):
        outcomes[i]["atlas"]["difference_q"] = qv
    categories = {}
    for i in families["A"]:
        out = outcomes[i]
        category = atlas_category(out)
        if out["testable"]:
            out["atlas"]["category"] = category
            m, hr = out["main"], out["atlas"]["hr"]
            if not out["atlas"]["comparable"] and hr > 1 and m["hi"] is not None:
                # The atlas's outcome is a subtype of the trials'. If only that subtype, a share s of
                # the whole, rose by the atlas's factor, the whole would rise to 1 + s (hr - 1); the
                # trials' upper limit allows s up to this.
                out["atlas"]["largest_share_consistent"] = max(0.0, (m["hi"] - 1) / (hr - 1))
        categories.setdefault(category, []).append(i)

    # The registered scans with the wider variance.
    half_of = A.halves(strata)
    scans = {"organ_systems": scan(strata, half_of, "soc", sorted(A.load("results/R4.json")["items"])),
             "preferred_terms": scan(strata, half_of, "pt", sorted(A.load("results/R5.json")["items"]))}
    for n in scans["preferred_terms"]["replicated_lower"] + scans["preferred_terms"]["replicated_higher"]:
        scans["preferred_terms"]["items"][n]["merged_rows"] = groups.merged_rows(n)

    # Preferred terms within bacterial infections and pneumonia.
    terms = {}
    for oid in REVISED["terms_listed_for"]:
        listed = []
        for term in terms_counted(strata, oid):
            rows = A.tables(strata, ("pt", term))
            if A.total_events(rows) >= A.P["min_events_term_scan"]:
                listed.append({"term": term, **wider(rows), "merged_rows": groups.merged_rows(term)})
        terms[oid] = sorted(listed, key=lambda x: -(x["glp1_affected"] + x["placebo_affected"]))

    drop = set(REVISED["coronavirus_pneumonia"])
    rows = tables_without(strata, "A21", drop)
    pneumonia_without = {"left_out": sorted(drop), **wider(rows), "leave_one_out": leave_one_out(rows)}

    trials = sorted({s["nct"] for s in strata})
    assert set(trials) == set(WINDOW), "every eligible trial has a window"
    window_counts = {"not_on_treatment": sum(1 for nct in trials if WINDOW[nct] != "on_treatment")}
    for nct in trials:
        window_counts[WINDOW[nct]] = window_counts.get(WINDOW[nct], 0) + 1
    window_participants = {}
    for s in strata:
        w = WINDOW[s["nct"]]
        window_participants[w] = window_participants.get(w, 0) + s["glp1"].at_risk + s["placebo"].at_risk

    A.write("V.json", {
        "outcomes": outcomes,
        "atlas_categories": {k: sorted(v) for k, v in categories.items()},
        "scans": scans,
        "terms_per_participant": arms,
        "preferred_terms": terms,
        "pneumonia_without_coronavirus": pneumonia_without,
        "window": {"trials": window_counts, "participants": window_participants, "coding": window_agreement()},
    })


def window_agreement():
    """How far the two window coders agreed: the share of trials they coded alike, and Cohen's kappa."""
    one, two = A.load("data/ae_window/coder1.json"), A.load("data/ae_window/coder2.json")
    assert [r["nct"] for r in one] == [r["nct"] for r in two]
    a, b = [r["window"] for r in one], [r["window"] for r in two]
    n = len(a)
    agreed = sum(x == y for x, y in zip(a, b))
    chance = sum((a.count(c) / n) * (b.count(c) / n) for c in set(a) | set(b))
    return {"trials": n, "agreed": agreed, "kappa": (agreed / n - chance) / (1 - chance),
            "third_reader": len(A.load("data/ae_window/third_reader.json"))}


def atlas_category(out):
    """Where the trials leave one of the atlas's associations, with both estimates' uncertainty:
    consistent (q < 0.05 in the atlas's direction, and no difference from it detected after
    correction, which is not a test of equivalence), same direction but different in size, opposite,
    different (nothing significant in the trials, but a difference from the atlas's estimate after
    correction), or inconclusive. Only outcomes significant in the trials can be consistent, so how
    many are is bounded by the trials' power."""
    if not out["testable"]:
        return "too_few_events"
    m, atlas = out["main"], out["atlas"]
    if A.significant(out["q"]):
        if (m["rr"] > 1) != (atlas["hr"] > 1):
            return "opposite"
        return "consistent" if not A.significant(atlas.get("difference_q")) else "same_direction_different_size"
    return "different" if A.significant(atlas.get("difference_q")) else "inconclusive"


if __name__ == "__main__":
    main()