"""Decimal and rational cross-checks; does not import either analysis implementation."""
from decimal import Decimal as D, 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())
source = json.loads((ROOT/'data/registry-source.json').read_text())
results = json.loads((ROOT/'results/full-results.json').read_text())
byid = {r['id']: r for r in rows}
selected = [r for r in rows if r['primary_pool']]

def calculate(records):
    with localcontext() as ctx:
        ctx.prec = 60
        ys, vs = [], []
        for r in records:
            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)
        ws = [1/v for v in vs]
        mean = sum(w*y for w,y in zip(ws,ys))/sum(ws)
        q = sum(w*(y-mean)**2 for w,y in zip(ws,ys))
        df = D(len(records)-1)
        tau = max(D(0),(q-df)/(sum(ws)-sum(w*w for w in ws)/sum(ws)))
        rw = [1/(v+tau) for v in vs]
        mu = sum(w*y for w,y in zip(rw,ys))/sum(rw)
        scale = max(D(1),sum(w*(y-mu)**2 for w,y in zip(rw,ys))/df)
        delta = (scale/sum(rw)).sqrt()*D(str(t.ppf(.975,int(df))))
        return {k:float(v) for k,v in dict(rr=mu.exp(),lower=(mu-delta).exp(),upper=(mu+delta).exp(),tau2=tau,fixed_rr=mean.exp()).items()}

def compare(a,b):
    for k,v in a.items():
        assert abs(v-b[k]) < 1e-10, (k,v,b[k])

compare(calculate(selected),results['primary'])
for r in selected:
    compare(calculate([x for x in selected if x['id'] != r['id']]),results['sensitivities']['leave_out_'+r['id']])
for k,v in source['counts'].items():
    assert byid['covidout'][k] == v
assert source['time_days'] == byid['covidout']['followup_days'] == 28
assert len(selected) == 4
assert sum(r['n_t']+r['n_c'] for r in selected) == 7714

tg = byid['together']
h,d = tg['hospital_only'],tg['death_only']
ut = {h['events_t']+d['events_t']-o for o in range(max(0,h['events_t']+d['events_t']-h['n_t']),min(h['events_t'],d['events_t'])+1)}
uc = {h['events_c']+d['events_c']-o for o in range(max(0,h['events_c']+d['events_c']-h['n_c']),min(h['events_c'],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 ut for c in uc}
for r in scenarios:
    compare(calculate(selected + [{**tg,'events_t':r['union_events_t'],'events_c':r['union_events_c']}]),r)
assert sum(r['lower'] <= 1 <= r['upper'] for r in scenarios) == results['together']['intervals_including_null']

co = byid['covidout']
assignments = json.loads((ROOT/'results/covidout-denominator-scenarios.json').read_text())
nt,nc = co['analysis_n_t'],co['analysis_n_c']
mt,mc = nt-co['n_t'],nc-co['n_c']
assert len(assignments) == (mt+1)*(mc+1)
assert {(r['additional_events_t'],r['additional_events_c']) for r in assignments} == {(a,c) for a in range(mt+1) for c in range(mc+1)}
for r in assignments:
    exact = Fraction(co['events_t']+r['additional_events_t'],nt)/Fraction(co['events_c']+r['additional_events_c'],nc)
    assert abs(float(exact)-r['rr']) < 1e-13
assert results['covidout_denominator']['first_treatment_events_reversing_direction'] == 2

# Reversing both study arms must reciprocate estimates and swap interval endpoints.
swapped = [{**r,'events_t':r['events_c'],'n_t':r['n_c'],'events_c':r['events_t'],'n_c':r['n_t']} for r in selected]
rev = calculate(swapped)
original = calculate(selected)
assert abs(rev['rr']*original['rr']-1) < 1e-12
assert abs(rev['lower']*original['upper']-1) < 1e-12
assert abs(rev['upper']*original['lower']-1) < 1e-12
report = {'decimal_primary_and_leave_one_out_match':True,'decimal_union_scenarios_match':len(scenarios),
          'rational_denominator_scenarios_match':len(assignments),'registry_extraction_matches':True,
          'reciprocal_symmetry_passed':True,'scope':'Independent arithmetic and extraction consistency, not independent patient-data validation.'}
(ROOT/'results/R2.json').write_text(json.dumps(report,indent=2)+'\n')
print(json.dumps(report))