import numpy as np,pandas as pd
from scipy import optimize
from load import STUDIES
rng=np.random.default_rng(20261101)
def suff(d):
    g=d.groupby(['cond','stim']).y.agg(['size','mean','var']).reset_index()
    g['var']=g['var'].fillna(0); ssw=((g['size']-1)*g['var']).sum()
    return g['cond'].values,g['size'].values.astype(float),g['mean'].values,ssw,len(d)
def reml(cond,n,m,ssw,N,starts=None):
    J=len(n);conds=np.unique(cond)
    def nll(p):
        se2=np.exp(p[0]);sv2=np.exp(p[1])
        w=n/(se2+n*sv2);ll=(N-J)*np.log(se2)+ssw/se2+np.sum(np.log(se2+n*sv2))
        for c in conds:
            k=cond==c;mu=np.sum(w[k]*m[k])/np.sum(w[k]);ll+=np.sum(w[k]*(m[k]-mu)**2)+np.log(np.sum(w[k]))
        return 0.5*ll
    se0=ssw/max(N-J,1);best=None
    for sv in (starts if starts else [se0*1e-4,se0*0.05,se0*0.5]):
        r=optimize.minimize(nll,[np.log(se0),np.log(max(sv,se0*1e-6))],method='L-BFGS-B',bounds=[(None,None),(np.log(se0)-25,np.log(se0)+8)])
        if best is None or r.fun<best.fun: best=r
    se2,sv2=np.exp(best.x)
    if sv2<1e-8*se2: sv2=0.0
    return sv2,se2
def ni(sv2,se2): return sv2/(sv2+se2)
def boot(cond,n,m,ssw,N,sv2,se2,B=1000):
    out=np.empty(B);J=len(n);conds=np.unique(cond)
    mu={c:np.average(m[cond==c],weights=n[cond==c]) for c in conds}
    for b in range(B):
        v=rng.normal(0,np.sqrt(sv2),J);mb=np.array([mu[c] for c in cond])+v+rng.normal(0,np.sqrt(se2/n))
        sswb=se2*rng.chisquare(N-J)
        a,e=reml(cond,n,mb,sswb,N,starts=[se2*1e-4,max(sv2,se2*1e-3)]);out[b]=ni(a,e)
    return out
def logit(p): p=np.clip(p,0.001,0.999);return np.log(p/(1-p))
def crossed(d):
    # balanced two-way ANOVA estimator, stimuli nested in condition, persons crossed
    W=d.pivot_table(index='subj',columns='stim',values='y',aggfunc='mean').dropna()
    I,J=W.shape;cond=d.drop_duplicates('stim').set_index('stim').cond.reindex(W.columns).values
    Y=W.values;gm=Y.mean();pm=Y.mean(1,keepdims=True);sm=Y.mean(0)
    cm=np.array([sm[cond==c].mean() for c in cond])
    E=Y-pm-sm[None,:]+gm;msr=(E**2).sum()/((I-1)*(J-1))
    C=len(np.unique(cond));msv=I*np.sum((sm-cm)**2)/(J-C)
    sv2=max(0,(msv-msr)/I);return ni(sv2,msr)
def spearman(d):
    vals=[];wts=[]
    for c,x in d.groupby('cond'):
        W=x.pivot_table(index='subj',columns='stim',values='y',aggfunc='mean').dropna()
        R=W.rank(axis=1).values;R=R-R.mean(1,keepdims=True);sd=np.sqrt((R**2).sum(1));ok=sd>0
        R=R[ok]/sd[ok,None];M=R@R.T;n=M.shape[0];vals.append((M.sum()-n)/(n*(n-1)));wts.append(W.shape[1])
    return np.average(vals,weights=wts)
import sys,json
rows=[]
for s,(cls,dsn,fn) in [(k,STUDIES[k]) for k in sys.argv[1:]]:
    d=fn();cond,n,m,ssw,N=suff(d);sv2,se2=reml(cond,n,m,ssw,N);bs=boot(cond,n,m,ssw,N,sv2,se2)
    r=dict(study=s,cls=cls,design=dsn,NI=ni(sv2,se2),lo=np.percentile(bs,2.5),hi=np.percentile(bs,97.5),
           logit=logit(ni(sv2,se2)),v_logit=np.var(logit(bs),ddof=1),zero=sv2==0,nstim=len(n))
    for c in ['save','del']:
        dc=d[d.cond==c]
        if dc.stim.nunique()>1:
            a,b,mm,w,NN=suff(dc);x1,x2=reml(a,b,mm,w,NN);r['NI_'+c]=ni(x1,x2)
    if dsn=='within': r['NI_crossed']=crossed(d);r['spearman']=spearman(d)
    rows.append(r);print(s,round(r['NI'],4),[round(r['lo'],4),round(r['hi'],4)],'zero' if r['zero'] else '',flush=True)
import os
for r in rows: open('res_'+r['study'].replace(' ','_').replace('&','and').replace('(','').replace(')','')+'.json','w').write(json.dumps({k:(bool(v) if isinstance(v,(np.bool_,bool)) else (float(v) if isinstance(v,(np.floating,float,int,np.integer)) and not isinstance(v,bool) else v)) for k,v in r.items()}))