"""Independent review of aggregate counting and MH estimates. No author code is imported.
Run: python independent_review.py BUNDLE_DIR OUTPUT_JSON
Bounds concern feasible unions of the posted adverse-event rows, not confidence intervals.
"""
import collections,hashlib,json,math,sys
from pathlib import Path
b=Path(sys.argv[1]);output=Path(sys.argv[2])
def load(p):return json.loads((b/p).read_text())
sc=load('data/screening.json');coding={x['term']:x for x in load('data/term_coding.json')};outs={x['id']:x for x in load('data/outcomes.json')};specs=load('data/analyses.json');molecules=load('data/molecules.json')
records={p.stem:json.loads(p.read_text()) for p in (b/'data/ctgov').glob('*.json')}
def norm(s):return ' '.join((s or '').split()).lower()
def age(s):
 if not s:return None
 v,u=s.split();return float(v)/{'years':1,'months':12,'weeks':365.25/7,'days':365.25}[u.lower()]
def at_risk(g,ae):
 if g.get('seriousNumAtRisk'):return g['seriousNumAtRisk']
 return max([x.get('numAtRisk') or 0 for e in ae.get('seriousEvents',[]) for x in e.get('stats',[]) if x['groupId']==g['id']]+[0]) or g.get('deathsNumAtRisk') or 0
def arm(ae,gids,outcome):
 gs={g['id']:g for g in ae['eventGroups']};terms={k for k,v in coding.items() if set(v['outcomes'])&set(outs[outcome]['codes'])};ns=total=lower=upper=0;details=[]
 for gid in gids:
  g=gs[gid];n=at_risk(g,ae);counts=collections.defaultdict(int);raw=[]
  for e in ae.get('seriousEvents',[]):
   t=norm(e.get('term'))
   if t not in terms:continue
   for x in e.get('stats',[]):
    if x['groupId']==gid and x.get('numAffected',0):
     raw.append((t,x['numAffected']));counts[coding[t]['canonical']]+=x['numAffected']
  s=sum(counts.values());lo=max((v for _,v in raw),default=0);hi=min(s,n,g['seriousNumAffected']) if g.get('seriousNumAffected') is not None else min(s,n)
  ns+=n;total+=s;lower+=lo;upper+=hi
  if raw:details.append({'group':gid,'title':g['title'],'n':n,'sum_terms':s,'union_lower':lo,'union_upper':hi,'any_serious':g.get('seriousNumAffected'),'terms':dict(counts),'raw':raw})
 return {'n':ns,'sum':min(total,ns),'lower':lower,'upper':upper,'groups':details}
strata=[];reused=[]
for row in sc:
 if not row['eligible']:continue
 rec=records[row['nct']];ae=rec['resultsSection']['adverseEventsModule'];el=rec['protocolSection']['eligibilityModule'];lo=age(el.get('minimumAge'));hi=age(el.get('maximumAge'))
 ag='pediatric' if hi is not None and hi<18 else 'older' if lo is not None and lo>=60 else 'adult' if lo is not None and lo>=18 else 'mixed'
 sm=rec['protocolSection']['statusModule'];date=(sm.get('primaryCompletionDateStruct') or sm.get('completionDateStruct') or {}).get('date','')[:7]
 used=collections.Counter(g for c in row['comparisons'] for k in ['antidepressant_groups','placebo_groups'] for g in c[k]);reused.extend({'nct':row['nct'],'group':k,'uses':v} for k,v in used.items() if v>1)
 for k,c in enumerate(row['comparisons'],1):
  a=arm(ae,c['antidepressant_groups'],'S1');p=arm(ae,c['placebo_groups'],'S1')
  if not a['n'] or not p['n']:continue
  strata.append({'nct':row['nct'],'id':f"{row['nct']}#{k}",'age_group':ag,'design':row['design'],'indication':row['indication'],'combination':bool(c.get('combination')),'classes':set(molecules[m] for m in c['molecules']),'date':date,'ae':ae,'comp':c})
def selected(s,sp):
 return all(k not in sp or s[k] in sp[k] for k in ['age_group','design','indication']) and (sp.get('combinations',False) or not s['combination']) and ('classes' not in sp or s['classes']<=set(sp['classes'])) and ('completed_after' not in sp or s['date']>sp['completed_after'])
def mh(rows,drug='sum',plac='sum'):
 R=S=V=0
 for a,p in rows:
  x,y=a[drug],p[plac];n,m=a['n'],p['n'];N=n+m
  R+=x*m/N;S+=y*n/N;V+=(n*m*(x+y)-x*y*N)/N**2
 if not R or not S:return None
 rr=R/S;se=math.sqrt(V/R/S);z=math.log(rr)/se
 return {'rr':rr,'lo':math.exp(math.log(rr)-1.95996398454*se),'hi':math.exp(math.log(rr)+1.95996398454*se),'p':math.erfc(abs(z)/math.sqrt(2))}
report={'records':len(records),'comparisons':len(strata),'trials':len({s['nct'] for s in strata}),'reused_groups':reused,'analyses':{},'code_matches_plan':(b/'code/analyze.py').read_bytes()==(b/'plan/code/analyze.py').read_bytes()}
for key in ['P1','P2','A6','A18','A19']:
 sp=next(x for x in specs if x['id']==key);ss=[s for s in strata if selected(s,sp['select'])];rows=[];amb=[];impossible=[]
 for s in ss:
  a=arm(s['ae'],s['comp']['antidepressant_groups'],sp['outcome']);p=arm(s['ae'],s['comp']['placebo_groups'],sp['outcome']);rows.append((a,p))
  for name,ar in [('drug',a),('placebo',p)]:
   for g in ar['groups']:
    if g['union_lower']!=g['union_upper']:amb.append({'nct':s['nct'],'arm':name,**g})
    if g['sum_terms']>g['union_upper']:impossible.append({'nct':s['nct'],'arm':name,**g})
 results={'trials':len({s['nct'] for s in ss}),'comparisons':len(rows),'counts':{armname:{v:sum(row[i][v] for row in rows) for v in ['n','sum','lower','upper']} for i,armname in enumerate(['drug','placebo'])},'mh_as_reported':mh(rows),'mh_min_union':mh(rows,'lower','upper'),'mh_max_union':mh(rows,'upper','lower'),'ambiguous_groups':amb,'sum_exceeds_unique_upper_groups':impossible,'pre2008_trials':len({s['nct'] for s in ss if s['date'] and s['date']<'2008-01'}),'single_zero_event_comparisons':sum((a['sum']==0)!=(p['sum']==0) for a,p in rows),'event_comparisons':sum(a['sum']+p['sum']>0 for a,p in rows)}
 report['analyses'][key]=results
checks=load('plan/ctgov_digests.json');report['snapshot_digests_matched']=sum('sha256:'+hashlib.sha256((b/'data/ctgov'/n).read_bytes()).hexdigest()==dg for n,dg in checks.items())
for name in ['screening.json','term_coding.json','outcomes.json','analyses.json','molecules.json','parameters.json','search.json']:
 report.setdefault('plan_data_matches',{})[name]=(b/'data'/name).read_bytes()==(b/'plan/data'/name).read_bytes()
output.write_text(json.dumps(report,indent=2)+'\n')
print(json.dumps({**{k:v for k,v in report.items() if k!='analyses'},'analyses':{k:{a:v for a,v in val.items() if a not in ['ambiguous_groups','sum_exceeds_unique_upper_groups']}|{'ambiguous_groups':len(val['ambiguous_groups']),'sum_exceeds_unique_upper_groups':len(val['sum_exceeds_unique_upper_groups'])} for k,val in report['analyses'].items()}},indent=2))