"""Recompute a targeted audit from frozen public trial-level counts.

All intervals are exploratory. No individual participant data are used.
"""
from copy import deepcopy
import json
from math import exp, log, sqrt
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm, t

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "results"
OUT.mkdir(exist_ok=True)


def effect(row):
    a, n, c, m = (row[k] for k in ("events_t", "n_t", "events_c", "n_c"))
    assert all(isinstance(v, int) for v in (a, n, c, m))
    assert 0 < a < n and 0 < c < m, "This analysis requires nonzero event cells."
    y = log((a / n) / (c / m))
    v = 1 / a - 1 / n + 1 / c - 1 / m
    z = norm.ppf(0.975)
    return dict(id=row["id"], label=row["label"], log_rr=y, variance=v,
                rr=exp(y), lower=exp(y-z*sqrt(v)), upper=exp(y+z*sqrt(v)),
                p_two_sided=2*norm.sf(abs(y)/sqrt(v)),
                risk_t=a/n, risk_c=c/m, risk_difference_pp=100*(a/n-c/m),
                events_t=a, n_t=n, events_c=c, n_c=m,
                endpoint=row["endpoint"], followup_days=row["followup_days"])


def pool(rows):
    es = [effect(r) for r in rows]
    y = np.array([r["log_rr"] for r in es])
    v = np.array([r["variance"] for r in es])
    k = len(es)
    assert k >= 2
    w = 1/v
    fixed = float(np.dot(w,y)/sum(w))
    q = float(np.dot(w,(y-fixed)**2))
    tau2 = max(0., (q-k+1)/(sum(w)-sum(w*w)/sum(w)))
    wr = 1/(v+tau2)
    mu = float(np.dot(wr,y)/sum(wr))
    hk_scale = max(1., float(np.dot(wr,(y-mu)**2)/(k-1)))
    se = sqrt(hk_scale/sum(wr))
    half = t.ppf(.975,k-1)*se
    z = norm.ppf(.975)
    return dict(k=k, n=sum(r["n_t"]+r["n_c"] for r in rows),
                n_t=sum(r["n_t"] for r in rows), n_c=sum(r["n_c"] for r in rows),
                events_t=sum(r["events_t"] for r in rows),
                events_c=sum(r["events_c"] for r in rows),
                rr=exp(mu), lower=exp(mu-half), upper=exp(mu+half),
                p_two_sided=float(2*t.sf(abs(mu/se),k-1)),
                tau2=float(tau2), q=q, i2_percent=max(0.,100*(q-k+1)/q) if q else 0.,
                hk_scale=hk_scale, fixed_rr=exp(fixed),
                fixed_lower=exp(fixed-z/sqrt(sum(w))),
                fixed_upper=exp(fixed+z/sqrt(sum(w))),
                fixed_p_two_sided=float(2*norm.sf(abs(fixed)*sqrt(sum(w)))),
                weights={r["id"]:float(wi/sum(wr)) for r,wi in zip(rows,wr)})


def main():
    rows = json.loads((ROOT/"data/trials.json").read_text())
    primary = [r for r in rows if r["primary_pool"]]
    effects = {r["id"]:effect(r) for r in rows}
    pooled = pool(primary)
    sensitivity = {"leave_out_"+r["id"]:pool([s for s in primary if s["id"]!=r["id"]])
                   for r in primary}
    sensitivity["blinded_only"] = pool([r for r in primary if r["blinded"]])
    sensitivity["add_14_day_covidout"] = pool(primary+[r for r in rows if r["id"]=="covidout"][0:1])
    for name, updates in [
        ("remove_accidental_death", {"events_t":6}),
        ("missing_placebo_all_nonevents", {"n_c":724}),
        ("missing_placebo_all_events", {"n_c":724,"events_c":4}),
    ]:
        altered = deepcopy(primary)
        next(r for r in altered if r["id"]=="activ600").update(updates)
        sensitivity[name]=pool(altered)
    bad = deepcopy(primary)
    p = next(r for r in bad if r["id"]=="principle")
    p.update({k:p["other_control"][k] for k in ("events_c","n_c")})
    invalid = effect(p)
    sensitivity["INVALID_all_principle_controls"] = pool(bad)
    together = next(r for r in rows if r["id"]=="together")
    hospital = {**together,**together["hospital_only"],"endpoint":"all-cause hospitalization alone"}
    effects["together_hospital_only"]=effect(hospital)
    # The union of hospitalizations and deaths is unidentified from its margins.
    # Frechet bounds: max(H,D) <= H union D <= min(N,H+D).
    h,d=together["hospital_only"],together["death_only"]
    amin,amax=max(h["events_t"],d["events_t"]),min(h["n_t"],h["events_t"]+d["events_t"])
    cmin,cmax=max(h["events_c"],d["events_c"]),min(h["n_c"],h["events_c"]+d["events_c"])
    bounds=dict(events_t_min=amin,events_t_max=amax,events_c_min=cmin,events_c_max=cmax,
                possible_crude_rr_min=(amin/h["n_t"])/(cmax/h["n_c"]),
                possible_crude_rr_max=(amax/h["n_t"])/(cmin/h["n_c"]))
    # Do not assume any death-hospitalization overlap. Enumerate every allowed
    # integer union and show the envelope of conditional meta-analysis outputs.
    union_scenarios=[]
    for a in range(amin,amax+1):
        for c in range(cmin,cmax+1):
            union={**together,"events_t":a,"events_c":c,
                   "endpoint":"hypothetical hospitalization-or-death union"}
            result=pool(primary+[union])
            union_scenarios.append({"union_events_t":a,"union_events_c":c,**result})
    (OUT/"together-union-scenarios.json").write_text(json.dumps(union_scenarios,indent=2)+"\n")
    bounds.update(scenarios=len(union_scenarios),pooled_n=union_scenarios[0]["n"],
                  pooled_rr_min=min(r["rr"] for r in union_scenarios),
                  pooled_rr_max=max(r["rr"] for r in union_scenarios),
                  lower_envelope=min(r["lower"] for r in union_scenarios),
                  upper_envelope=max(r["upper"] for r in union_scenarios),
                  intervals_including_null=sum(r["lower"]<=1<=r["upper"] for r in union_scenarios))
    full=dict(effects=effects,primary=pooled,sensitivities=sensitivity,
              invalid_principle=invalid,together_union_identification_bounds=bounds)
    (OUT/"full-results.json").write_text(json.dumps(full,indent=2,allow_nan=False)+"\n")
    # Display precision is part of the declared result; full precision remains above.
    rounded=lambda x:round(float(x),3)
    display=lambda k,v:float(f"{v:.6g}") if "p_two_sided" in k else rounded(v) if isinstance(v,float) else v
    summary=dict(source_comparisons=len(rows),source_platforms=len({r["platform"] for r in rows}),
                 source_analysis_records=sum(r["analysis_n_t"]+r["analysis_n_c"] for r in rows),
                 interval_level_percent=95,
                 primary={k:display(k,v) for k,v in pooled.items() if k!="weights"},
                 principle_concurrent={k:rounded(v) for k,v in effects["principle"].items() if k in ["rr","lower","upper","risk_difference_pp"]},
                 principle_nonconcurrent={k:rounded(v) for k,v in invalid.items() if k in ["rr","lower","upper","risk_difference_pp"]},
                 controls_added=3256-1806,control_events_added=144-27,
                 bias_ratio=rounded(invalid["rr"]/effects["principle"]["rr"]),
                 together_union_rr_min=rounded(bounds["possible_crude_rr_min"]),
                 together_union_rr_max=rounded(bounds["possible_crude_rr_max"]))
    summary["union_sensitivity"]={k:display(k,v) for k,v in bounds.items()}
    summary["trials"]={r["id"]:{k:display(k,v) for k,v in effects[r["id"]].items()
                               if k not in ["variance","log_rr","label","id"]} for r in rows}
    summary["sensitivities"]={name:{k:display(k,v) for k,v in val.items() if k!="weights"}
                              for name,val in sensitivity.items()}
    (OUT/"R1.json").write_text(json.dumps(summary,indent=2,allow_nan=False)+"\n")
    make_plot(effects,pooled,invalid)
    print(json.dumps(summary,indent=2))


def make_plot(effects,pooled,invalid):
    fig,axes=plt.subplots(2,1,figsize=(10.5,7.5),gridspec_kw={"height_ratios":[2,1]},layout="constrained")
    color="#176d75"
    points=[effects[k] for k in ("principle","activ400","activ600")]+[{**pooled,"label":"Exploratory pooled estimate (modified HK)"}]
    for i,row in enumerate(points):
        axes[0].errorbar(row["rr"],i,xerr=[[row["rr"]-row["lower"]],[row["upper"]-row["rr"]]],
                         fmt="D" if i==3 else "o",color=color,capsize=4)
    axes[0].set_yticks(range(4),[r["label"] for r in points]);axes[0].invert_yaxis()
    axes[0].set_title("Reported 28-day hospitalization or death",loc="left",weight="bold")
    axes[0].set_xlabel("Risk ratio and 95% CI • pooled interval accounts for few comparisons")
    for i,row in enumerate([effects["principle"],invalid]):
        axes[1].errorbar(row["rr"],i,xerr=[[row["rr"]-row["lower"]],[row["upper"]-row["rr"]]],
                         fmt="o",color=color if i==0 else "#b65d39",capsize=4)
    axes[1].set_yticks([0,1],["Concurrent eligible controls","All controls (invalid crude comparison)"]);axes[1].invert_yaxis()
    axes[1].set_title("PRINCIPLE: control choice changes the apparent effect",loc="left",weight="bold")
    axes[1].set_xlabel("Risk ratio and 95% CI • lower than 1 favors ivermectin")
    for ax in axes:
        ax.set_xscale("log");ax.axvline(1,color="#66717a",lw=1,ls="--");ax.grid(axis="x",alpha=.2)
        ax.spines[["top","right"]].set_visible(False)
    axes[0].set_xticks([.5,1,2,5,10],["0.5","1","2","5","10"])
    axes[1].set_xticks([.25,.5,1,2],["0.25","0.5","1","2"])
    fig.suptitle("Ivermectin: public randomized-trial results",fontsize=17,weight="bold")
    fig.savefig(OUT/"trial-audit.png",dpi=160)
    plt.close(fig)


if __name__=="__main__":
    main()