import sys,json,math,collections
sys.path.insert(0,'code')
import analyze as a
strata=a.build_strata(); definitions={x['id']:x for x in a.ANALYSES};out={}
def bounds(s,outcome,armkey):
 comp=a.SCREENING[s['nct']]['comparisons'][int(s['id'].split('#')[1])-1]
 rec=a.record(s['nct']); mod=rec['resultsSection']['adverseEventsModule'];groups={g['id']:g for g in mod['eventGroups']}
 low=high=0;details=[]
 for gid in comp['antidepressant_groups' if armkey=='drug' else 'placebo_groups']:
  vals=[st.get('numAffected',0) or 0 for row in mod.get('seriousEvents',[]) if a.norm(row.get('term')) in a.outcome_terms(outcome) for st in row.get('stats',[]) if st['groupId']==gid]
  l=max(vals,default=0);h=sum(vals);cap=groups[gid].get('seriousNumAffected')
  if cap is not None:
   h=min(h,cap)
   outside=sum(st.get('numAffected',0) or 0 for row in mod.get('seriousEvents',[]) if a.norm(row.get('term')) not in a.outcome_terms(outcome) for st in row.get('stats',[]) if st['groupId']==gid)
   l=max(l,cap-outside)
  assert l<=h,(s['id'],gid,l,h)
  low+=l;high+=h
  if len([v for v in vals if v])>1:details.append({'group':gid,'counts':vals,'lower':l,'upper':h})
 return low,high,details
for ident in ['P1','P2','A6','A18','A19']:
 spec=definitions[ident];sel=a.select(strata,spec['select']); rows=a.tables(sel,spec['outcome']);loRR=[];hiRR=[];bothlow=[];details=[];tot=[0,0,0,0];original=a.mantel_haenszel(rows)
 for s,x,n,y,m in rows:
  dl,du,dd=bounds(s,spec['outcome'],'drug');pl,pu,pd=bounds(s,spec['outcome'],'placebo');tot=[v+w for v,w in zip(tot,[dl,du,pl,pu])]
  loRR.append((s,dl,n,pu,m));hiRR.append((s,du,n,pl,m));bothlow.append((s,dl,n,pl,m))
  if dd or pd:details.append({'comparison':s['id'],'drug':dd,'placebo':pd})
 out[ident]={'original':original,'union_count_bounds':tot,'lowest_rr_feasible':a.mantel_haenszel(loRR),'highest_rr_feasible':a.mantel_haenszel(hiRR),'maximum_overlap_both_arms':a.mantel_haenszel(bothlow),'ambiguous_groups':details}
sel=a.select(strata,definitions['P1']['select']);missing=[];old=[];repeat=[];seen=set();med=[]
for s in sel:
 rec=a.record(s['nct']);comp=a.SCREENING[s['nct']]['comparisons'][int(s['id'].split('#')[1])-1];groups={g['id']:g for g in rec['resultsSection']['adverseEventsModule']['eventGroups']}
 for gid in comp['antidepressant_groups']+comp['placebo_groups']:
  key=(s['nct'],gid)
  if key in seen:repeat.append(key)
  seen.add(key)
  if groups[gid].get('seriousNumAffected') is None:missing.append(key)
 if s['completion'] and s['completion']<'2008':old.append({'id':s['id'],'date':s['completion']})
 measures=rec.get('resultsSection',{}).get('baselineCharacteristicsModule',{}).get('measures',[])
 ages=[m for m in measures if a.re.sub('[^a-z]','',m.get('title','').lower())=='agecontinuous']
 if ages and not any(m.get('paramType')=='MEAN' for m in ages):med.append(s['nct'])
out['audits']={'missing_any_SAE':missing,'repeated_event_group_ids':repeat,'pre2008_completion':old,'median_age_only':sorted(set(med))}
# Independent contrast of two uncertain estimates: illustrative normal approximation, assumes independent samples and comparable estimands, neither established.
for ident,prior,lo,hi in [('P1',.87,.58,1.29),('P2',.79,.64,.98)]:
 est=out[ident]['original'];priorse=math.log(hi/lo)/(2*1.959964);z=(math.log(est['rr']/prior))/math.sqrt(est['se']**2+priorse**2)
 out[ident]['illustrative_two_estimate_comparison']={'z':z,'p':math.erfc(abs(z)/math.sqrt(2)),'assumptions':'independent normal log estimates; common estimand, not established'}
a.os.makedirs('results',exist_ok=True)
json.dump(out,open('results/reviewer_audit.json','w'),indent=2)
print(json.dumps({k:{x:v for x,v in y.items() if x not in ['ambiguous_groups']} for k,y in out.items() if k!='audits'},indent=2))
print(json.dumps(out['audits']))