"""Rounds the registered outputs (results/R1.json to R6.json) and the analyses added after review
(results/V.json, code/revised.py) 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; what code/revised.py found is under "revised". 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"], "trials": m.get("trials"), "trials_with_events": m.get("trials_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 slug(name):
    """A key for a named item that a placeholder can name: lower case, every run of other characters one underscore."""
    out, gap = [], False
    for ch in name.lower():
        if ch.isascii() and ch.isalnum():
            out.append(ch)
            gap = False
        elif not gap:
            out.append("_")
            gap = True
    return "".join(out).strip("_")


def keyed(names, value):
    """{slug: value(name)}, refusing two names with one slug."""
    out = {}
    for name in names:
        key = slug(name)
        assert key not in out, (name, key)
        out[key] = {"name": name, **value(name)}
    return out


def scan(r):
    def item(name):
        it = r["items"][name]
        return {"all": estimate(it["all"]), "discovery": estimate(it["discovery"]), "discovery_q": pval(it["discovery_q"]),
                "replication": estimate(it["replication"]), "replicated": it["replicated"]}
    items = keyed(r["discovered"], item)
    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")),
        "revised": revised(load("V.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 percent(x):
    return None if x is None else round(100 * x)


def revised_outcome(o):
    out = {"name": o["name"], "counts": o["counts"], "preferred_terms": o["preferred_terms"], "all": estimate(o["main"]),
           "q": pval(o.get("q")), "robust": o.get("robust")}
    if not o["testable"]:
        return out
    if o.get("leave_one_out"):
        lo = o["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"])}
    w = o["window"]
    out["window"] = {k: estimate(w[k]) for k in ("on_treatment", "in_trial", "calendar", "unstated", "not_on_treatment")}
    out["window"]["interaction_p"] = pval(w["interaction_p"])
    if "overlap" in o:
        ov = o["overlap"]
        t, shift = ov["tipping_ratio"], ov["within_trial_shift"]
        out["overlap"] = {k: {"rr": ratio(ov[k]["rr"]), "lo": ratio(ov[k]["lo"]), "hi": ratio(ov[k]["hi"]), "p": pval(ov[k]["p"])}
                          for k in ("adjusted", "pooled")}
        out["overlap"].update({"tipping_percent": None if t is None else percent(t - 1),
                               "within_trial_percent": None if shift is None else round(100 * (shift - 1), 1),
                               "bounds": ov["bounds"]})
    if "atlas" in o:
        a = o["atlas"]
        out["atlas"] = {"hr": a["hr"], "lo": a["lo"], "hi": a["hi"], "ratio": ratio(a["ratio"]), "ratio_lo": ratio(a["ratio_lo"]),
                        "ratio_hi": ratio(a["ratio_hi"]), "category": a["category"],
                        "difference_p": pval(a["difference_p"]), "difference_q": pval(a.get("difference_q")),
                        "comparable": a["comparable"]}
        if "largest_share_consistent" in a:
            out["atlas"]["largest_share_percent"] = percent(a["largest_share_consistent"])
    return out


def revised_scan(r):
    def item(name):
        it = r["items"][name]
        out = {"all": estimate(it["all"]), "discovery": estimate(it["discovery"]), "discovery_q": pval(it["discovery_q"]),
               "replication": estimate(it["replication"]), "replicated": it["replicated"]}
        if "merged_rows" in it:
            out["merged_rows"] = it["merged_rows"]
        return out
    return {"tested": r["tested"], "discovered": r["discovered"], "replicated_lower": r["replicated_lower"],
            "replicated_higher": r["replicated_higher"], "discovered_count": len(r["discovered"]),
            "replicated_lower_count": len(r["replicated_lower"]), "replicated_higher_count": len(r["replicated_higher"]),
            "items": keyed(r["discovered"], item)}


def revised(v):
    cats = v["atlas_categories"]
    p = v["pneumonia_without_coronavirus"]
    return {
        "outcomes": {k: revised_outcome(o) for k, o in v["outcomes"].items()},
        "atlas": {"counts": {c: len(ids) for c, ids in cats.items()}, "lists": cats},
        "scans": {k: revised_scan(r) for k, r in v["scans"].items()},
        "terms_per_participant": {arm: round(a["terms_per_participant"], 2) for arm, a in v["terms_per_participant"].items()},
        "preferred_terms": {oid: keyed([t["term"] for t in terms], lambda n, ts=terms: {
                                **estimate(next(t for t in ts if t["term"] == n)),
                                "merged_rows": next(t for t in ts if t["term"] == n)["merged_rows"]})
                            for oid, terms in v["preferred_terms"].items()},
        "pneumonia_without_coronavirus": {**estimate(p), "leave_one_out_p_max": pval(p["leave_one_out"]["p_max"])},
        "window": {**v["window"], "coding": {**v["window"]["coding"], "kappa": round(v["window"]["coding"]["kappa"], 2)}},
    }


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()