import json,sys,math,csv,statistics
from pathlib import Path
from scipy.stats import beta,binom,chi2,t,norm
from scipy.integrate import quad
p=Path(sys.argv[1]);d=json.loads((p/'results/associations.json').read_text());rows=list(d.values());n=len(rows)
def bh(v):
 order=sorted(range(len(v)),key=v.__getitem__);ans=[0]*len(v);prev=1
 for i in range(len(v)-1,-1,-1):
  j=order[i];prev=min(prev,v[j]*len(v)/(i+1));ans[j]=prev
 return ans
def share(k,n):return {'k':k,'n':n,'share':round(k/n,3),'ci':[round(beta.ppf(.025,k,n-k+1),3) if k else 0,round(beta.ppf(.975,k+1,n-k),3) if k<n else 1]}
def med(x):
 x=sorted(x);lo=int(binom.ppf(.025,len(x),.5));return {'median':round(statistics.median(x),3),'ci':[round(x[lo-1],3),round(x[len(x)-lo],3)]}
def power(r,alpha):
 f=r['replication'];df=f['df'];df=math.inf if df is None else df;nc=r['published_b'][0]/f['se'];critical=t.ppf(1-alpha/2,df)
 if math.isinf(df):return norm.cdf(-critical-nc)+norm.sf(critical-nc)
 return quad(lambda v: (norm.sf(critical*math.sqrt(v/df)-nc)+norm.cdf(-critical*math.sqrt(v/df)-nc))*chi2.pdf(v,df),0,math.inf,epsabs=1e-10)[0]
q=bh([r['replication']['p'] for r in rows]);rep=[math.copysign(1,r['replication']['b'])==math.copysign(1,r['published_b'][0]) and v<.05 for r,v in zip(rows,q)];inf=[power(r,.05)>=.8 for r in rows];infb=[power(r,.05/n)>=.8 for r in rows]
repz=[];rept=[]
for r in rows:
 f=r['replication'];pub=r['published_b'];se=(pub[2]-pub[1])/(2*norm.ppf(.975));z=(f['b']-pub[0])/math.sqrt(f['se']**2+se**2);df=f['df'] if f['df'] is not None else math.inf;repz.append(2*norm.sf(abs(z)));rept.append(2*t.sf(abs(z),df))
head=[x for r in rows for x in r['departures'] if x['affects_headline']]
coding=list(csv.DictReader((p/'data/departure_coding.csv').open()));agree=sum(r['coder_1']==r['coder_2'] for r in coding);cats=['paper','data','unresolved'];chance=sum(sum(r['coder_1']==c for r in coding)*sum(r['coder_2']==c for r in coding) for c in cats)/len(coding)**2
out={'n':n,'replicated':share(sum(rep),n),'informative':sum(inf),'replicated_informative':share(sum(a and b for a,b in zip(rep,inf)),sum(inf)),'informative_bonferroni':sum(infb),'replicated_informative_bonferroni':sum(a and b for a,b in zip(rep,infb)),'ratio':med([r['replication']['b']/r['published_b'][0] for r in rows]),'ratio_informative':med([r['replication']['b']/r['published_b'][0] for r,i in zip(rows,inf) if i]),'reproduced':sum(r['published']['low']<=r['original']['estimate']<=r['published']['high'] for r in rows),'difference_z':sum(v<.05 for v in bh(repz)),'difference_t':sum(v<.05 for v in bh(rept)),'headline_departures':len(head),'by_evidence':{c:sum(x['evidence']==c for x in head) for c in cats},'coding_agree':agree,'coding_kappa':round((agree/len(coding)-chance)/(1-chance),2),'maximum_BH_discrepancy':max(abs(a-r['q']) for a,r in zip(q,rows))}
out.update({'ratio_original':med([r['original']['b']/r['published_b'][0] for r in rows]),'ratio_own':med([r['replication']['b']/r['harmonized']['b'] for r in rows]),'inside_published_ci':sum(r['published']['low']<=r['replication']['estimate']<=r['published']['high'] for r in rows),'same_sign_nominal':sum(r['same_sign'] and r['replication']['p']<.05 for r in rows),'reversed_corrected':sum(not r['same_sign'] and v<.05 for r,v in zip(rows,q)),'papers_with_departure':sum(any(x['affects_headline'] for x in r['departures']) for r in rows),'papers_with_paper_evidence':sum(any(x['affects_headline'] and x['evidence']=='paper' for x in r['departures']) for r in rows),'papers_with_data_only_evidence':sum(not any(x['affects_headline'] and x['evidence']=='paper' for x in r['departures']) and any(x['affects_headline'] and x['evidence']=='data' for x in r['departures']) for r in rows),'heterogeneous':sum(v<.05 for v in bh([r['heterogeneity_p'] for r in rows]))})
print(json.dumps(out,indent=2,default=lambda x:x.item()))
