"""Writes the registered outputs (results/R1.json to R5.json and comparisons.csv) and the
correction's (results/R6.json) as the paper and the claims give them.

Writes results/S.json, every value as the paper shows it: risk ratios and intervals to two
decimals, p- and q-values as the style guide writes them (two significant digits, or
"< 0.001"), percentages per year of age to one decimal, and counts with thousands separators.
A ratio that two decimals would round to 1.00, or a p-value that two significant digits would
round to 0.050, gets the digits it takes to show which side of 1 or of 0.05 it falls on. It
computes nothing new beyond crude rates and the median of the trials' mean ages. Also writes
results/analyses.csv, every registered analysis in one table, for readers.

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

import csv
import json
import math
import os
import statistics
import sys

ROOT = sys.argv[1] if len(sys.argv) > 1 else "."
RES = os.path.join(ROOT, "results")
ALPHA = 0.05
MINUS = "−"


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


def signed(text):
    return MINUS + text[1:] if text.startswith("-") else text


def ratio(x):
    """Two decimals, or up to four when two would round a ratio other than 1 to 1.00."""
    if x is None:
        return None
    for digits in (2, 3, 4):
        if round(x, digits) != 1.0 or x == 1.0:
            break
    return f"{x:.{digits}f}"


def pval(p):
    """Two significant digits, keeping a trailing zero (0.10), or "< 0.001"; three when two
    would round a p-value other than 0.05 to 0.050."""
    if p is None:
        return None
    if p < 0.001:
        return "< 0.001"
    digits = max(2, 1 - math.floor(math.log10(p)))
    if f"{p:.{digits}f}" == f"{ALPHA:.{digits}f}" and p != ALPHA:
        digits += 1
    return f"{p:.{digits}f}"


def count(n):
    return None if n is None else f"{n:,}"


def one(x):
    """One decimal, with a minus sign."""
    return None if x is None else signed(f"{x:.1f}")


def span(lo, hi):
    """A range as the paper writes it, "43 to 46", or one value when both ends agree."""
    return lo if lo == hi else f"{lo} to {hi}"


# ---------------------------------------------------------------------------
# The registered outputs
# ---------------------------------------------------------------------------


def estimate(m):
    return {"rr": ratio(m.get("rr")), "lo": ratio(m.get("lo")), "hi": ratio(m.get("hi")), "p": pval(m.get("p")),
            "drug_affected": count(m["drug_affected"]), "drug_at_risk": count(m["drug_at_risk"]),
            "placebo_affected": count(m["placebo_affected"]), "placebo_at_risk": count(m["placebo_at_risk"]),
            "participants": count(m["drug_at_risk"] + m["placebo_at_risk"]),
            "comparisons": count(m["comparisons"]), "comparisons_with_events": count(m["comparisons_with_events"])}


def analysis(e):
    out = {"name": e["name"], "affected": count(e["events"]), "testable": e["testable"], "all": estimate(e["main"]),
           "p": pval(e["p_primary"])}
    if "q" in e:
        out["q"] = pval(e["q"])
    if "p_holm" in e:
        out["p_holm"] = pval(e["p_holm"])
    if e.get("against_priors"):
        out["against"] = e["against_priors"]
    if e.get("power"):
        out["power"] = {k: f"{v:.2f}" for k, v in e["power"].items()}
    pe = e["peto"]
    out["peto"] = {"or": ratio(pe["or"]), "lo": ratio(pe["lo"]), "hi": ratio(pe["hi"]), "p": pval(pe["p"])}
    if not e["testable"]:
        return out
    re = e["random_effects"]
    out["random_effects"] = {"rr": ratio(re["rr"]), "lo": ratio(re["lo"]), "hi": ratio(re["hi"]), "p": pval(re["p"]),
                             "i2": None if re.get("i2") is None else f"{re['i2']:.2f}"}
    out["conservative_p"] = pval(e["conservative"]["p"])
    out["overlap_bound"] = f"{e['overlap_bound']:.2f}"
    out["half1"] = estimate(e["half1"])
    out["half2"] = estimate(e["half2"])
    if "leave_one_out" in e:
        lo = e["leave_one_out"]
        out["leave_one_out"] = {"rr_min": ratio(lo["rr_min"]), "rr_max": ratio(lo["rr_max"]), "lo_min": ratio(lo["lo_min"]),
                                "hi_max": ratio(lo["hi_max"]), "p_min": pval(lo["p_min"]), "p_max": pval(lo["p_max"])}
    if "dominant_term" in e:
        out["dominant_term"] = {"term": e["dominant_term"]["term"].lower(), **estimate(e["dominant_term"])}
    if "robust" in e:
        out["robust"] = e["robust"]
    return out


# ---------------------------------------------------------------------------
# The correction's outputs
# ---------------------------------------------------------------------------


def interval(m):
    return {"rr": ratio(m["rr"]), "lo": ratio(m["lo"]), "hi": ratio(m["hi"]), "p": pval(m["p"])}


def bounds(b):
    """An analysis at every count the posted tables allow."""
    every = b["every"]
    out = {"name": b["name"], "testable": b["testable"], "ambiguous_arms": count(len(b["ambiguous"])),
           "drug_least": count(b["drug_least"]), "drug_most": count(b["drug_most"]),
           "placebo_least": count(b["placebo_least"]), "placebo_most": count(b["placebo_most"]),
           "drug_at_risk": count(b["drug_at_risk"]), "placebo_at_risk": count(b["placebo_at_risk"]),
           "drug": span(count(b["drug_least"]), count(b["drug_most"])),
           "placebo": span(count(b["placebo_least"]), count(b["placebo_most"])),
           "configurations": count(every["configurations"]), "significant": count(every["significant"])}
    if math.isfinite(every["rr_min"]):
        out.update({"rr_min": ratio(every["rr_min"]), "rr_max": ratio(every["rr_max"]), "lo_min": ratio(every["lo_min"]),
                    "lo_max": ratio(every["lo_max"]), "hi_min": ratio(every["hi_min"]), "hi_max": ratio(every["hi_max"]),
                    "p_min": pval(every["p_min"]), "p_max": pval(every["p_max"])})
        out["rr"] = span(out["rr_min"], out["rr_max"])
        out["p"] = span(out["p_min"], out["p_max"])
    for name in ("lowest", "highest", "shared", "separate"):
        out[name] = {"drug_affected": count(b[name]["drug_affected"]), "placebo_affected": count(b[name]["placebo_affected"]),
                     **interval(b[name])}
    for key in ("holm_at_p_max", "holm_at_p_min", "holm_shared", "holm_separate", "q_at_p_max", "q_at_p_min"):
        if b.get(key) is not None:
            out[key] = pval(b[key])
    if "q_at_p_max" in out:
        out["q"] = span(out["q_at_p_min"], out["q_at_p_max"])
    return out


def corrected(r6):
    out = {"bounds": {k: bounds(v) for k, v in r6["bounds"].items()},
           "single_term": {k: {**bounds({"name": v["term"], **v}), "term": v["term"].lower()} for k, v in r6["single_term"].items()},
           "estimators": {}, "age": {}}
    for k, v in r6["estimators"].items():
        dl = v["registered_random_effects"]
        out["estimators"][k] = {
            "comparisons_with_events": count(v["comparisons_with_events"]), "with_a_zero_cell": count(v["with_a_zero_cell"]),
            "dl": {**interval(dl), "tau2": f"{dl['tau2']:.2f}", "i2": f"{dl['i2']:.2f}"},
            "inverse_variance_corrected": interval(v["inverse_variance_corrected"]),
        }
        for name in ("registered", "lowest", "highest"):
            for model in ("random", "common"):
                m = v[name][model]
                out["estimators"][k][f"{name}_{model}"] = {**interval(m), "tau": f"{m['tau']:.2f}"}
    for k, v in r6["age"].items():
        out["age"][k] = {"configurations": count(v["configurations"]), "pct_min": one(v["pct_min"]), "pct_max": one(v["pct_max"]),
                         "p_min": pval(v["p_min"]), "p_max": pval(v["p_max"]),
                         "comparisons_with_events": count(v["comparisons_with_events"]),
                         "median_age_comparisons": count(len(v["median_age_trials"]))}
    out["checks"] = {k: bounds({"name": k, **v}) for k, v in r6["checks"].items()}
    out["resolved"] = {k: bounds(v) for k, v in r6["resolved"].items()}
    out["resolution"] = {k: count(v) for k, v in r6["resolution"].items()}
    out["independence"] = {}
    for k, v in r6["independence"].items():
        if k == "pediatric":
            out["independence"][k] = {"comparisons": count(v["comparisons"]), "completed_after_fda_2004": count(v["completed_after_fda_2004"]),
                                      "first_completion": v["first_completion"][:4]}
            continue
        out["independence"][k] = {part: {**bounds({"name": part, **b}), "comparisons": count(b["comparisons"]), "trials": count(b["trials"]),
                                         "participants": count(b["participants"]), "resolved": bounds({"name": part, **b["resolved"]})}
                                  for part, b in v.items()}
    # Whether each earlier estimate lies inside the 95% interval at every count the tables allow,
    # at none, or at some.
    out["against"] = {}
    for k, v in list(r6["bounds"].items()) + [(f"{k}_neither", v["neither"]) for k, v in r6["independence"].items() if k != "pediatric"]:
        every = v["every"]
        if every["contains"]:
            out["against"][k] = {name: "every" if n == every["configurations"] else "none" if n == 0 else "some"
                                 for name, n in every["contains"].items()}
    return out


def main():
    r1, r2, r3, r4, r5, r6 = (load(f"R{i}.json") for i in range(1, 7))
    with open(os.path.join(RES, "comparisons.csv"), encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    adult = [float(r["mean_age"]) for r in rows if r["design"] == "acute" and r["age_group"] in ("adult", "older") and r["mean_age"]]
    pediatric = [float(r["mean_age"]) for r in rows if r["design"] == "acute" and r["age_group"] == "pediatric" and r["mean_age"]]
    a1 = r1["analyses"]
    s = {
        "trials": {
            "found": count(r1["found"]), "screened": count(r1["screened"]), "eligible": count(r1["eligible_trials"]),
            "comparisons": count(r1["all"]["comparisons"]), "participants": count(r1["all"]["participants"]),
            "drug": count(r1["all"]["drug_participants"]), "placebo": count(r1["all"]["placebo_participants"]),
            "pediatric_trials": count(a1["A19"]["trials"]), "pediatric_participants": count(a1["A19"]["participants"]),
            "adult_trials": count(a1["P1"]["trials"]), "adult_comparisons": count(a1["P1"]["comparisons"]),
            "adult_participants": count(a1["P1"]["participants"]), "adult_drug": count(a1["P1"]["drug_participants"]),
            "adult_placebo": count(a1["P1"]["placebo_participants"]),
            "adult_depressive_participants": count(a1["P1"]["by_indication"].get("depressive", 0)),
            "adult_depressive_trials": count(a1["A6"]["trials"]),
            "adult_median_mean_age": one(statistics.median(adult)) if adult else None,
            "pediatric_median_mean_age": one(statistics.median(pediatric)) if pediatric else None,
            "adult_with_mean_age": count(len(adult)),
        },
        "primary": {k: analysis(v) for k, v in r2.items()},
        "secondary": {k: analysis(v) for k, v in r3.items()},
        "age": {k: {"name": v["name"], "comparisons": count(v["comparisons"]),
                    "pct_per_year": one(v.get("pct_per_year")), "pct_lo": one(v.get("pct_lo")), "pct_hi": one(v.get("pct_hi")),
                    "p": pval(v.get("p")), "age_min": one(v.get("age_min")), "age_max": one(v.get("age_max"))} for k, v in r4.items()},
        "checks": {k: {"name": v["name"], "affected": count(v["events"]), "all": estimate(v["main"]), "p": pval(v["p_primary"]),
                       "as_expected": v.get("as_expected")} for k, v in r5.items()},
        "corrected": corrected(r6),
    }
    s["secondary_significant"] = sorted(k for k, v in r3.items() if v.get("q") is not None and v["q"] < ALPHA)
    # Crude rates per 10,000 participants, for display beside the risk ratios: the registered
    # counts, and the bounds the correction puts on them.
    for k, v in r2.items():
        m = v["main"]
        b = r6["bounds"][k]
        per = lambda n, d: one(1e4 * n / d)
        s["primary"][k]["per_10000"] = {"drug": per(m["drug_affected"], m["drug_at_risk"]), "placebo": per(m["placebo_affected"], m["placebo_at_risk"]),
                                        "drug_least": per(b["drug_least"], b["drug_at_risk"]), "drug_most": per(b["drug_most"], b["drug_at_risk"]),
                                        "placebo_least": per(b["placebo_least"], b["placebo_at_risk"]),
                                        "placebo_most": per(b["placebo_most"], b["placebo_at_risk"])}
    # The registered earlier estimates and the post-2017 cutoff, copied from data/ for display.
    with open(os.path.join(ROOT, "data", "parameters.json"), encoding="utf-8") as f:
        s["priors"] = {k: f"{v:.2f}" for k, v in json.load(f)["priors"].items()}
    # The earlier estimates as their papers give them (data/earlier.json, added in the correction).
    with open(os.path.join(ROOT, "data", "earlier.json"), encoding="utf-8") as f:
        s["earlier"] = {k: {"value": f"{v['value']:.2f}", "lo": f"{v['lo']:.2f}", "hi": f"{v['hi']:.2f}"}
                        for k, v in json.load(f)["estimates"].items()}
    with open(os.path.join(ROOT, "data", "analyses.json"), encoding="utf-8") as f:
        s["completed_after_year"] = next(a for a in json.load(f) if a["id"] == "A10")["select"]["completed_after"][:4]
    with open(os.path.join(RES, "S.json"), "w", encoding="utf-8") as f:
        json.dump(s, f, indent=1, sort_keys=True, ensure_ascii=False)
        f.write("\n")
    with open(os.path.join(RES, "analyses.csv"), "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["id", "analysis", "drug_affected", "drug_at_risk", "placebo_affected", "placebo_at_risk", "rr", "lo", "hi", "p", "adjusted_p",
                    "testable", "drug_least", "drug_most", "placebo_least", "placebo_most", "rr_lowest", "rr_highest", "p_min", "p_max",
                    "adjusted_p_at_p_max"])
        for fam, key, at_max in ((r2, "p_holm", "holm_at_p_max"), (r3, "q", "q_at_p_max")):
            for k, v in fam.items():
                m, b = v["main"], r6["bounds"][k]
                e = b["every"]
                fin = math.isfinite(e["rr_min"])
                w.writerow([k, v["name"], m["drug_affected"], m["drug_at_risk"], m["placebo_affected"], m["placebo_at_risk"],
                            m["rr"], m["lo"], m["hi"], m["p"], v.get(key), v["testable"], b["drug_least"], b["drug_most"],
                            b["placebo_least"], b["placebo_most"], e["rr_min"] if fin else None, e["rr_max"] if fin else None,
                            e["p_min"] if fin else None, e["p_max"] if fin else None, b.get(at_max)])


if __name__ == "__main__":
    main()