"""Group-preserving identification bounds; stdlib only; no clinical term classifier."""
import csv,gzip,hashlib,io,json,sys
from collections import Counter
from fractions import Fraction as F
from pathlib import Path
R=Path(__file__).resolve().parents[1]

def norm(s):return ' '.join(s.split()).lower()
def age(s):
 if not s:return None
 value,unit=s.split()
 return F(value)/{'Year':1,'Years':1,'Month':12,'Months':12,'Week':F('52.1775'),'Weeks':F('52.1775'),'Day':F('365.25'),'Days':F('365.25')}[unit]
def agegroup(rec):
 el=rec['protocolSection']['eligibilityModule'];lo=age(el.get('minimumAge'));hi=age(el.get('maximumAge'))
 if hi is not None and hi<18:return 'pediatric'
 if lo is not None and lo>=60:return 'older'
 if lo is not None and lo>=18:return 'adult'
 return 'mixed'
def bounds(n,selected,other,s=None):
 assert n>=0 and all(isinstance(x,int) and 0<=x<=n for x in selected+other)
 lo=max(selected,default=0);up=min(n,sum(selected));valid=s is not None and max(selected+other,default=0)<=s<=min(n,sum(selected+other))
 return lo,up,max(lo,s-sum(other)) if valid else lo,min(up,s) if valid else up,valid

def main():
 src=json.loads(gzip.decompress((R/'data/frozen-source.json.gz').read_bytes()));idx=json.loads((R/'data/source-index.json').read_text())
 for name,raw in src.items():assert idx['files'][name]=='sha256:'+hashlib.sha256(raw.encode()).hexdigest(),name
 load=lambda n:json.loads(src[n])
 screen=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={x['id']:x for x in load('data/analyses.json') if x['id'] in ['P1','P2','A18','A19']}
 source_rows=list(csv.DictReader(io.StringIO(src['results/analyses.csv'])))
 groups=[];comparisons=[];diagnostics=[]
 for row in sorted(screen,key=lambda x:x['nct']):
  if not row['eligible'] or row['design']!='acute':continue
  nct=row['nct'];rec=load('data/ctgov/'+nct+'.json');ag=agegroup(rec);ae=rec['resultsSection']['adverseEventsModule'];events=ae.get('seriousEvents',[]);eg={x['id']:x for x in ae['eventGroups']}
  for ident,spec in specs.items():
   if ag not in spec['select']['age_group']:continue
   codes=set(outs[spec['outcome']]['codes'])
   for k,comp in enumerate(row['comparisons'],1):
    if comp.get('combination'):continue
    arms={}
    for arm,key in [('drug','antidepressant_groups'),('placebo','placebo_groups')]:
     aa=[]
     for gid in comp[key]:
      g=eg[gid];stats=[(i,e,t) for i,e in enumerate(events) for t in e.get('stats',[]) if t['groupId']==gid]
      n=g.get('seriousNumAtRisk') or max((t.get('numAtRisk') or 0 for _,_,t in stats),default=0) or g.get('deathsNumAtRisk') or 0
      terms=[]
      for i,e,t in stats:
       term=norm(e['term']);assert term in coding,term
       a=t.get('numAffected') or 0
       terms.append({'row':i,'term':e['term'],'organ':e.get('organSystem'),'selected':bool(codes.intersection(coding[term]['outcomes'])),'count':a,'at_risk':t.get('numAtRisk')})
      sel=[x['count'] for x in terms if x['selected']];other=[x['count'] for x in terms if not x['selected']]
      lo,up,rl,ru,valid=bounds(n,sel,other,g.get('seriousNumAffected'))
      mismatches=[x for x in terms if x['at_risk'] is not None and x['at_risk']!=n and x['count']>0]
      if mismatches:rl,ru,valid=lo,up,False
      assert not any(x['selected'] and x['count'] and x['at_risk'] not in (None,n) for x in terms),'selected population mismatch'
      key_terms=[(x['term'],x['organ']) for x in terms];dup=len(key_terms)-len(set(key_terms))
      item={'analysis':ident,'comparison':f'{nct}#{k}','nct':nct,'group':gid,'arm':arm,'n':n,'selected_counts':sel,'other_counts':other,'terms':terms,'s':g.get('seriousNumAffected'),'lower':lo,'upper':up,'refined_lower':rl,'refined_upper':ru,'refinement_valid':valid,'positive_denominator_mismatches':len(mismatches),'duplicate_term_organ_rows':dup}
      aa.append(item)
     arms[arm]=aa
    n1=sum(x['n'] for x in arms['drug']);n0=sum(x['n'] for x in arms['placebo'])
    if not n1 or not n0:
     diagnostics.append({'analysis':ident,'comparison':f'{nct}#{k}','excluded':'zero denominator'});continue
    for a in arms.values():groups.extend(a)
    cr={'analysis':ident,'comparison':f'{nct}#{k}','n1':n1,'n0':n0}
    for arm,gg in arms.items():
     for metric in ['lower','upper','refined_lower','refined_upper']:cr[arm+'_'+metric]=sum(g[metric] for g in gg)
     cr[arm+'_sum_capped_after_pool']=min(sum(sum(g['selected_counts']) for g in gg),sum(g['n'] for g in gg))
    comparisons.append(cr)
 summary={}
 for ident in specs:
  cc=[x for x in comparisons if x['analysis']==ident];gg=[x for x in groups if x['analysis']==ident]
  reuse=Counter((x['nct'],x['group']) for x in gg)
  # Reused same-arm groups would preserve monotonicity, but opposite-arm reuse would not.
  roles={}
  for x in gg:roles.setdefault((x['nct'],x['group']),set()).add(x['arm'])
  assert all(len(x)==1 for x in roles.values()),'cross-arm reuse needs a separate optimizer'

  def mh(dm,pm):
   a=sum((F(x['drug_'+dm]*x['n0'],x['n1']+x['n0']) for x in cc),F(0));c=sum((F(x['placebo_'+pm]*x['n1'],x['n1']+x['n0']) for x in cc),F(0))
   assert c>0
   return a/c
  vals={'lower':mh('lower','upper'),'upper':mh('upper','lower'),'refined_lower':mh('refined_lower','refined_upper'),'refined_upper':mh('refined_upper','refined_lower'),'capped_sum':mh('sum_capped_after_pool','sum_capped_after_pool')}
  summary[ident]={'trials':len({x['comparison'].split('#')[0] for x in cc}),'comparisons':len(cc),'groups':len(gg),'reused_groups':sum(v>1 for v in reuse.values()),'groups_with_width':sum(x['upper']>x['lower'] for x in gg),'groups_refinement_unavailable':sum(not x['refinement_valid'] for x in gg),'groups_denominator_mismatch':sum(x['positive_denominator_mismatches']>0 for x in gg),'duplicate_term_organ_rows':sum(x['duplicate_term_organ_rows'] for x in gg),'drug_at_risk':sum(x['n1'] for x in cc),'placebo_at_risk':sum(x['n0'] for x in cc),**{k:float(v) for k,v in vals.items()},'exact':{k:f'{v.numerator}/{v.denominator}' for k,v in vals.items()}}
  for arm in ['drug','placebo']:
   for metric in ['lower','upper','sum_capped_after_pool']:summary[ident][arm+'_'+metric]=sum(x[arm+'_'+metric] for x in cc)
  summary[ident]['display']={k:round(float(v),4) for k,v in vals.items()}
  # Confirm fixed published source counts and point statistic, not its inferential intervals.
  old=next(x for x in source_rows if x.get('id',x.get('analysis'))==ident)
  for arm in ['drug','placebo']:
   assert summary[ident][arm+'_sum_capped_after_pool']==int(old[arm+'_affected'])
   assert summary[ident][arm+'_at_risk']==int(old[arm+'_at_risk'])
  assert abs(summary[ident]['capped_sum']-float(old['rr']))<0.005
  summary[ident]['source_row']=old
 out=R/'results';out.mkdir(exist_ok=True)
 for name,obj in [('groups.json',groups),('comparisons.json',comparisons),('diagnostics.json',diagnostics),('R.json',summary)]:
  if name=='groups.json':(out/(name+'.gz')).write_bytes(gzip.compress(json.dumps(obj,sort_keys=True,separators=(',',':')).encode(),mtime=0))
  else:(out/name).write_text(json.dumps(obj,indent=2,sort_keys=True)+'\n')
 print(json.dumps(summary,indent=2))
if __name__=='__main__':main()