"""Rounds the registered outputs (results/R1.json to R5.json and comparisons.csv) for the paper
and the claims.

Writes results/S.json: 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, counts as counted, and the median of the trials' mean ages. It computes nothing new.
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")


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


def ratio(x):
    return None if x is None else round(x, 2)


def pval(p):
    """Two significant digits, keeping a trailing zero (0.10), or "< 0.001"."""
    if p is None:
        return None
    if p < 0.001:
        return "< 0.001"
    digits = max(2, 1 - math.floor(math.log10(p)))
    return f"{p:.{digits}f}"


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": m["drug_affected"], "drug_at_risk": m["drug_at_risk"],
            "placebo_affected": m["placebo_affected"], "placebo_at_risk": m["placebo_at_risk"],
            "comparisons": m["comparisons"], "comparisons_with_events": m["comparisons_with_events"]}


def analysis(e):
    out = {"name": e["name"], "affected": 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: round(v, 2) 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 round(re["i2"], 2)}
    out["conservative_p"] = pval(e["conservative"]["p"])
    out["overlap_bound"] = round(e["overlap_bound"], 2)
    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"], **estimate(e["dominant_term"])}
    if "robust" in e:
        out["robust"] = e["robust"]
    return out


def main():
    r1, r2, r3, r4, r5 = (load(f"R{i}.json") for i in range(1, 6))
    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"]]
    s = {
        "trials": {
            "found": r1["found"], "screened": r1["screened"], "eligible": r1["eligible_trials"],
            "comparisons": r1["all"]["comparisons"], "participants": r1["all"]["participants"],
            "drug": r1["all"]["drug_participants"], "placebo": r1["all"]["placebo_participants"],
            "pediatric_trials": r1["analyses"]["A19"]["trials"], "pediatric_participants": r1["analyses"]["A19"]["participants"],
            "adult_trials": r1["analyses"]["P1"]["trials"], "adult_comparisons": r1["analyses"]["P1"]["comparisons"],
            "adult_participants": r1["analyses"]["P1"]["participants"], "adult_drug": r1["analyses"]["P1"]["drug_participants"],
            "adult_placebo": r1["analyses"]["P1"]["placebo_participants"],
            "adult_depressive_participants": r1["analyses"]["P1"]["by_indication"].get("depressive", 0),
            "adult_depressive_trials": r1["analyses"]["A6"]["trials"],
            "adult_median_mean_age": round(statistics.median(adult), 1) if adult else None,
            "pediatric_median_mean_age": round(statistics.median(pediatric), 1) if pediatric else None,
            "adult_with_mean_age": 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": v["comparisons"],
                    "pct_per_year": None if v.get("slope") is None else round(v["pct_per_year"], 1),
                    "pct_lo": None if v.get("slope") is None else round(v["pct_lo"], 1),
                    "pct_hi": None if v.get("slope") is None else round(v["pct_hi"], 1),
                    "p": pval(v.get("p")), "age_min": None if v.get("age_min") is None else round(v["age_min"], 1),
                    "age_max": None if v.get("age_max") is None else round(v["age_max"], 1)} for k, v in r4.items()},
        "checks": {k: {"name": v["name"], "affected": v["events"], "all": estimate(v["main"]), "p": pval(v["p_primary"]),
                       "as_expected": v.get("as_expected")} for k, v in r5.items()},
    }
    s["secondary_significant"] = sorted(k for k, v in r3.items() if v.get("q") is not None and v["q"] < 0.05)
    # Crude rates per 10,000 participants, for display beside the risk ratios.
    for v in s["primary"].values():
        a = v["all"]
        v["per_10000"] = {"drug": round(1e4 * a["drug_affected"] / a["drug_at_risk"], 1), "placebo": round(1e4 * a["placebo_affected"] / a["placebo_at_risk"], 1)}
    # 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"] = json.load(f)["priors"]
    with open(os.path.join(ROOT, "data", "analyses.json"), encoding="utf-8") as f:
        s["completed_after_year"] = int(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"])
        for fam, key in ((s["primary"], "p_holm"), (s["secondary"], "q")):
            for k, v in fam.items():
                a = v["all"]
                w.writerow([k, v["name"], a["drug_affected"], a["drug_at_risk"], a["placebo_affected"], a["placebo_at_risk"],
                            a["rr"], a["lo"], a["hi"], v["p"], v.get(key), v["testable"]])


if __name__ == "__main__":
    main()