import convergence as C
import numpy as np,pandas as pd,scipy.sparse as sp,json,os
from scipy.spatial import cKDTree
meta=pd.read_csv(C.NQ+'/data-neuroquery_version-1_metadata.tsv.gz',sep='\t');co=pd.read_csv(C.NQ+'/data-neuroquery_version-1_coordinates.tsv.gz',sep='\t');ids=meta.id.to_numpy();index={x:i for i,x in enumerate(ids)};co=co[co.id.isin(index)];ri=co.id.map(index).to_numpy();pts=co[['x','y','z']].to_numpy();has=np.zeros(len(ids),bool);has[np.unique(ri)]=True
act={}
for n,v in C.regions().items():
 d,_=cKDTree(v).query(pts);a=np.zeros(len(ids),bool);a[np.unique(ri[d<=10])]=True;act[n]=a[has]
X=sp.load_npz(C.NQ+'/data-neuroquery_version-1_vocab-neuroquery7547_source-body_type-count_features.npz').tocsc()[has];vocab=[x.rstrip('\n').split('\t')[0] for x in open(C.NQ+'/data-neuroquery_version-1_vocab-neuroquery7547_vocabulary.txt')];a=act['anterior_insula'];b=act['posterior_insula'];out={'n':len(a),'anterior':int(a.sum()),'posterior':int(b.sum()),'both':int((a&b).sum()),'posterior_also_anterior_fraction':float((a&b).sum()/b.sum()),'terms':{}}
for name,words in C.TARGETS.items():
 m=np.zeros(len(a),bool)
 for word in words:m|=(X[:,vocab.index(word)].toarray().ravel()>=2)
 p1=a[m].mean();p0=a[~m].mean();se=np.sqrt(p1*(1-p1)/m.sum()+p0*(1-p0)/(~m).sum());out['terms'][name]={'n':int(m.sum()),'anterior_with':int(a[m].sum()),'anterior_without':int(a[~m].sum()),'risk_difference':float(p1-p0),'iid_wald95':[float(p1-p0-1.96*se),float(p1-p0+1.96*se)],'anterior_minus_posterior_with':float(a[m].mean()-b[m].mean())}
os.makedirs('results',exist_ok=True);json.dump(out,open('results/reviewer.json','w'),indent=2);print(json.dumps(out))