"""Independent Decimal implementation and exact checks, not importing analysis."""
from decimal import Decimal, localcontext
from fractions import Fraction
import json
from pathlib import Path
from scipy.stats import t

ROOT=Path(__file__).resolve().parents[1]
rows=json.loads((ROOT/"data/trials.json").read_text())
full=json.loads((ROOT/"results/full-results.json").read_text())
cells=json.loads((ROOT/"data/source-count-audit.json").read_text())
byid={r["id"]:r for r in rows}
for item in cells:
    for key,value in item["numeric_counts"].items():
        assert byid[item["id"]][key]==value

with localcontext() as ctx:
    ctx.prec=60
    D=Decimal
    selected=[r for r in rows if r["primary_pool"]]
    ys=[];vs=[]
    for r in selected:
        a,n,c,m=[D(r[k]) for k in ("events_t","n_t","events_c","n_c")]
        ys.append((a/n).ln()-(c/m).ln())
        vs.append(1/a-1/n+1/c-1/m)
    weights=[1/v for v in vs]
    weighted_mean=sum(w*y for w,y in zip(weights,ys))/sum(weights)
    Q=sum(w*(y-weighted_mean)**2 for w,y in zip(weights,ys))
    df=D(len(ys)-1)
    heterogeneity=max(D(0),(Q-df)/(sum(weights)-sum(w*w for w in weights)/sum(weights)))
    rw=[1/(v+heterogeneity) for v in vs]
    mean=sum(w*y for w,y in zip(rw,ys))/sum(rw)
    scale=max(D(1),sum(w*(y-mean)**2 for w,y in zip(rw,ys))/df)
    error=(scale/sum(rw)).sqrt()*D(str(t.ppf(.975,int(df))))
    expected={"rr":mean.exp(),"lower":(mean-error).exp(),"upper":(mean+error).exp(),
              "fixed_rr":weighted_mean.exp(),"tau2":heterogeneity}
    for key,value in expected.items():
        assert abs(float(value)-full["primary"][key])<1e-11,(key,value)

p=byid["principle"]
concurrent=Fraction(p["events_t"],p["n_t"])/Fraction(p["events_c"],p["n_c"])
older=p["other_control"]
allcontrols=Fraction(p["events_t"],p["n_t"])/Fraction(older["events_c"],older["n_c"])
assert abs(float(concurrent)-full["effects"]["principle"]["rr"])<1e-14
assert abs(float(allcontrols)-full["invalid_principle"]["rr"])<1e-14

tr=byid["together"]
h,d=tr["hospital_only"],tr["death_only"]
# Enumerate the overlap, rather than the union, to check all allowable pairs.
unions_t={h["events_t"]+d["events_t"]-overlap for overlap in range(d["events_t"]+1)}
unions_c={h["events_c"]+d["events_c"]-overlap for overlap in range(d["events_c"]+1)}
scenarios=json.loads((ROOT/"results/together-union-scenarios.json").read_text())
assert {(r["union_events_t"],r["union_events_c"]) for r in scenarios}=={(a,c) for a in unions_t for c in unions_c}
assert min(unions_t)==79 and max(unions_t)==100
assert min(unions_c)==95 and max(unions_c)==119
assert all(r["lower"]<=1<=r["upper"] for r in scenarios)
result={"source_counts_match":True,"source_count_cells":sum(len(i["numeric_counts"]) for i in cells),
        "decimal_meta_analysis_match":True,"principle_exact_ratios_match":True,
        "all_overlap_pairs_covered":True,"overlap_pairs":len(scenarios),
        "all_overlap_intervals_include_one":True,
        "scope":"Independent arithmetic and extraction consistency; not independent participant-data validation."}
(ROOT/"results/R2.json").write_text(json.dumps(result,indent=2)+"\n")
print(json.dumps(result))
