"""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)})