"""Independent aggregate counts and exploratory uncertainty for the review."""
import ast
import json
from pathlib import Path
import nibabel as nib
import numpy as np
import pandas as pd
from scipy import sparse, stats
from scipy.spatial import cKDTree

root = Path('data/neuroquery')
def file(s):
    return next(root.glob('*'+s))
meta = pd.read_csv(file('metadata.tsv.gz'),sep='\t')
coords = pd.read_csv(file('coordinates.tsv.gz'),sep='\t')
vocab = [l.split('\t')[0] for l in file('vocabulary.txt').read_text().splitlines()]
features = sparse.load_npz(file('features.npz')).tocsc()
tree = ast.parse(Path('code/convergence.py').read_text())
decl = {n.targets[0].id:ast.literal_eval(n.value) for n in tree.body
        if isinstance(n,ast.Assign) and isinstance(n.targets[0],ast.Name)
        and n.targets[0].id in ('TARGETS','SYMPTOMS')}
terms = decl['TARGETS'] | decl['SYMPTOMS']
row = {identifier:i for i,identifier in enumerate(meta.id)}
coords = coords[coords.id.isin(row)]
ix = coords.id.map(row).to_numpy()
has = np.zeros(len(meta),bool); has[np.unique(ix)] = True
points = coords[['x','y','z']].to_numpy()
img = nib.load('data/HarvardOxford-cort-maxprob-thr25-2mm.nii.gz')
xyz = nib.affines.apply_affine(img.affine,np.argwhere(np.asarray(img.dataobj)==2))
dist = [cKDTree(xyz[xyz[:,1]>=0]).query(points)[0],cKDTree(xyz[xyz[:,1]<0]).query(points)[0]]
out = {'study_count':int(has.sum()),'radii':{},
       'scope':'Post hoc descriptive sensitivity. Wilson-Newcombe marginal intervals and approximate Welch intervals for the paired-region contrast assume independent articles; they ignore publication dependence and misclassification. They are not corrected for multiplicity or tests of clinical mechanisms.'}
def wilson(k,n):
    z=stats.norm.ppf(.975); p=k/n; den=1+z*z/n
    center=(p+z*z/(2*n))/den
    half=z*np.sqrt(p*(1-p)/n+z*z/(4*n*n))/den
    return center-half,center+half
def difference(a,b):
    p,q=a.mean(),b.mean(); la,ua=wilson(a.sum(),len(a));lb,ub=wilson(b.sum(),len(b))
    d=p-q
    return {'difference_pp':float(d*100),'ci95_pp':[float((d-np.sqrt((p-la)**2+(ub-q)**2))*100),float((d+np.sqrt((ua-p)**2+(q-lb)**2))*100)]}
for radius in (0,5,10):
    active=[]
    for distance in dist:
        hit=np.zeros(len(meta),bool);hit[np.unique(ix[distance<=radius])]=True
        active.append(hit[has])
    a,p=active; N=len(a)
    r={'joint_study_counts':{'both':int((a&p).sum()),'anterior_only':int((a&~p).sum()),'posterior_only':int((p&~a).sum()),'neither':int((~p&~a).sum())},'terms':{}}
    for name,words in terms.items():
        used=np.asarray((features[:,[vocab.index(w) for w in words]]>=2).sum(axis=1)).ravel()>0
        used=used[has]
        entry={'term_studies':int(used.sum()),'other_studies':int((~used).sum())}
        for label,h in zip(('anterior','posterior'),active):
            h1,h0=h[used],h[~used];pooled=h.mean()
            z=(h1.mean()-h0.mean())/np.sqrt(pooled*(1-pooled)*(1/len(h1)+1/len(h0)))
            entry[label]={'reported_with':int(h1.sum()),'reported_without':int(h0.sum()),'z':float(z),**difference(h1,h0)}
        d=a.astype(float)-p.astype(float);d1,d0=d[used],d[~used]
        v1,v0=d1.var(ddof=1)/len(d1),d0.var(ddof=1)/len(d0)
        se=np.sqrt(v1+v0);df=(v1+v0)**2/(v1*v1/(len(d1)-1)+v0*v0/(len(d0)-1))
        contrast=d1.mean()-d0.mean();half=stats.t.ppf(.975,df)*se
        entry['paired_region_difference_of_differences']={'estimate_pp':float(contrast*100),'ci95_pp':[float((contrast-half)*100),float((contrast+half)*100)],'df':float(df)}
        r['terms'][name]=entry
    out['radii'][str(radius)]=r
Path('results/reviewer_audit.json').write_text(json.dumps(out,indent=2)+'\n')
print(json.dumps(out,indent=2))