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

Writes results/S.json: risk ratios and their intervals to two decimals, p- and q-values as the
style guide writes them (two significant digits, or "< 0.001"), counts as counted, and the
lists the claims name. It computes nothing new.

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

import csv
import json
import math
import os
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):
    """A Mantel-Haenszel estimate as the paper shows it."""
    return {
        "rr": ratio(m.get("rr")), "lo": ratio(m.get("lo")), "hi": ratio(m.get("hi")), "p": pval(m.get("p")),
        "glp1_affected": m["glp1_affected"], "glp1_at_risk": m["glp1_at_risk"],
        "placebo_affected": m["placebo_affected"], "placebo_at_risk": m["placebo_at_risk"],
        "comparisons_with_events": m["strata_with_events"],
    }


def outcome(e):
    out = {"name": e["name"], "affected": e["events"], "testable": e["testable"], "all": estimate(e["main"])}
    if "va" in e:
        out["va_hr"] = e["va"]["hr"]
        out["va_category"] = e["va_category"]
        out["power_for_va_hr"] = round(e["power_for_va_hr"], 2)
    if not e["testable"]:
        return out
    out["q"] = pval(e["q"])
    out["robust"] = e["robust"]
    out["half1"] = estimate(e["half1"])
    out["half2"] = estimate(e["half2"])
    out["conservative_p"] = pval(e["conservative"]["p"])
    out["overlap_bound"] = round(e["overlap_bound"], 2)
    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)}
    pe = e["peto"]
    out["peto"] = {"or": ratio(pe["or"]), "lo": ratio(pe["lo"]), "hi": ratio(pe["hi"]), "p": pval(pe["p"])}
    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"]), "p_max": pval(lo["p_max"]),
                                "p_max_left_out": lo["p_max_left_out"]}
    if "dominant_term" in e:
        out["dominant_term"] = {"term": e["dominant_term"]["term"], **estimate(e["dominant_term"])}
    out["selective_only"] = estimate(e["selective_only"])
    out["type2_diabetes"] = estimate(e["type2_diabetes"])
    out["without_diabetes"] = estimate(e["without_diabetes"])
    out["population_interaction_p"] = pval(e.get("population_interaction_p"))
    return out


def scan(r):
    items = {}
    for name in r["discovered"]:
        it = r["items"][name]
        items[name] = {"all": estimate(it["all"]), "discovery": estimate(it["discovery"]), "discovery_q": pval(it["discovery_q"]),
                       "replication": estimate(it["replication"]), "replicated": it["replicated"]}
    return {"tested": r["tested"], "discovered": r["discovered"], "replicated": r["replicated"],
            "discovered_count": len(r["discovered"]), "replicated_count": len(r["replicated"]),
            "replicated_all_lower": all(r["items"][n]["all"]["rr"] < 1 for n in r["replicated"]), "items": items}


def main():
    r1, r2, r3, r6 = load("R1.json"), load("R2.json"), load("R3.json"), load("R6.json")
    outcomes = {k: outcome(v) for k, v in {**r2, **r3}.items()}
    by_category = {}
    for k, v in r2.items():
        by_category.setdefault(v["va_category"], []).append(k)
    counts = {c: len(by_category.get(c, [])) for c in ("confirmed", "opposite", "excludes_va_estimate", "inconclusive", "too_few_events")}
    informative = sorted(k for k, v in r2.items() if v.get("informative"))
    s = {
        "trials": {
            "found": r1["found"], "screened": r1["screened"], "eligible": r1["eligible_trials"], "comparisons": r1["comparisons"],
            "participants": r1["participants"], "glp1": r1["glp1_participants"], "placebo": r1["placebo_participants"],
            "half1": r1["half1_trials"], "half2": r1["half2_trials"], "overcount_factor": round(r1["overcount_factor"], 2),
            "type2_diabetes": r1["participants_by_population"].get("type2_diabetes", 0),
            "without_diabetes": r1["participants_by_population"].get("overweight_obesity", 0),
        },
        "outcomes": outcomes,
        "va": {
            "outcomes": r6["outcomes"], "testable": r6["testable"], "informative": r6["informative"],
            "counts": counts, "lists": {c: sorted(v) for c, v in by_category.items()},
            "informative_list": informative,
            "informative_counts": {c: sum(1 for k in informative if r2[k]["va_category"] == c) for c in counts},
            "robust": r6["robust"],
        },
        "family_b_significant": sorted(k for k, v in r3.items() if v["testable"] and v["q"] is not None and v["q"] < 0.05),
        "family_b_robust": sorted(k for k, v in r3.items() if v.get("robust")),
        "organ_systems": scan(load("R4.json")),
        "preferred_terms": scan(load("R5.json")),
    }
    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")
    write_table(outcomes)


def write_table(outcomes):
    """Every family A and B outcome in one table, for readers."""
    cols = ["id", "outcome", "glp1_affected", "glp1_at_risk", "placebo_affected", "placebo_at_risk", "rr", "lo", "hi", "p", "q",
            "robust", "random_effects_rr", "random_effects_p", "conservative_p", "half1_rr", "half2_rr", "va_hr", "va_category", "power_for_va_hr"]
    with open(os.path.join(RES, "outcomes.csv"), "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(cols)
        for k, o in outcomes.items():
            a = o["all"]
            w.writerow([k, o["name"], a["glp1_affected"], a["glp1_at_risk"], a["placebo_affected"], a["placebo_at_risk"], a["rr"], a["lo"], a["hi"], a["p"],
                        o.get("q"), o.get("robust"), o.get("random_effects", {}).get("rr"), o.get("random_effects", {}).get("p"), o.get("conservative_p"),
                        o.get("half1", {}).get("rr"), o.get("half2", {}).get("rr"), o.get("va_hr"), o.get("va_category"), o.get("power_for_va_hr")])


if __name__ == "__main__":
    main()