"""Independent checks: exact population enumeration, constructive witnesses, source joins."""
import csv,gzip,io,itertools,json,runpy
from collections import defaultdict
from fractions import Fraction
from pathlib import Path
R=Path(__file__).resolve().parents[1]

def assign(counts,u,start=0):
 assert max(counts,default=0)<=u<=sum(counts) if counts else u==0
 off=0;out=[]
 for count in counts:
  out.append({start+(off+j)%u for j in range(count)} if count else set());off+=count
 assert [len(s) for s in out]==counts
 assert len(set().union(*out))==u
 return out

def enumerate_small():
 marginal=defaultdict(set);total=defaultdict(set);populations=0
 for n in range(6):
  for m in range(1,5):
   for population in itertools.combinations_with_replacement(range(1<<m),n):
    counts=tuple(sum(bool(mask&(1<<j)) for mask in population) for j in range(m));S=sum(mask!=0 for mask in population);populations+=1
    for k in range(m+1):
     u=sum(bool(mask&((1<<k)-1)) for mask in population)
     marginal[n,k,counts].add(u);total[n,k,counts,S].add(u)
 for (n,k,counts),values in marginal.items():
  lo=max(counts[:k],default=0);hi=min(n,sum(counts[:k]))
  assert values==set(range(lo,hi+1))
 for (n,k,counts,S),values in total.items():
  lo=max(max(counts[:k],default=0),S-sum(counts[k:]));hi=min(S,sum(counts[:k]))
  assert values==set(range(lo,hi+1))
 return {'enumerated_populations':populations,'marginal_cases':len(marginal),'total_constrained_cases':len(total)}

def main():
 report=enumerate_small()
 f=runpy.run_path(str(R/'code/analyze.py'))['bounds']
 assert f(10,[3,4],[2],None)==(4,7,4,7,False)
 assert f(10,[3,4],[2],9)==(4,7,7,7,True)
 assert f(10,[3,4],[2],10)==(4,7,4,7,False)
 assert f(0,[],[],0)==(0,0,0,0,True)
 for args in [(-1,[],[],None),(3,[4],[],None),(3,[-1],[],None)]:
  try:f(*args)
  except AssertionError:pass
  else:raise AssertionError('invalid counts accepted')
 report['guard_fixtures']=7
 groups=json.loads(gzip.decompress((R/'results/groups.json.gz').read_bytes()));summ=json.loads((R/'results/R.json').read_text());src=json.loads(gzip.decompress((R/'data/frozen-source.json.gz').read_bytes()))
 coding={x['term']:set(x['outcomes']) for x in json.loads(src['data/term_coding.json'])};comps=list(csv.DictReader(io.StringIO(src['results/comparisons.csv'])));screen=json.loads(src['data/screening.json']);combo={f'{s["nct"]}#{i}':bool(c.get('combination')) for s in screen for i,c in enumerate(s.get('comparisons',[]),1)}
 witness=0;by=defaultdict(lambda:defaultdict(dict))
 for g in groups:
  rec=json.loads(src['data/ctgov/'+g['nct']+'.json']);ev=rec['resultsSection']['adverseEventsModule'];raw=[];expected_sel=[];expected_other=[]
  codes={'S1'} if g['analysis'] in ['P1','A18'] else {'S1','S2'}
  for e in ev.get('seriousEvents',[]):
   for stat in e.get('stats',[]):
    if stat['groupId']==g['group']:
     term=' '.join(e['term'].split()).lower();(expected_sel if coding[term]&codes else expected_other).append(stat.get('numAffected') or 0)
  assert expected_sel==g['selected_counts'] and expected_other==g['other_counts']
  for metric in ['lower','upper','refined_lower','refined_upper']:
   u=g[metric];a=assign(g['selected_counts'],u)
   if metric.startswith('refined') and g['refinement_valid']:
    v=max(max(g['other_counts'],default=0),g['s']-u)
    b=assign(g['other_counts'],v,g['s']-v)
    assert len(set().union(*(a+b)))==g['s']
    assert all(0<=x<g['n'] for part in a+b for x in part)
   witness+=1
  by[g['analysis']][g['comparison']].setdefault(g['arm'],[]).append(g)
 for ident,cc in by.items():
  wanted={x['comparison'] for x in comps if x['design']=='acute' and x['age_group'] in (['adult','older'] if ident in ['P1','P2'] else ['pediatric']) and not combo[x['comparison']]}
  assert set(cc)==wanted
  assert len(cc)==summ[ident]['comparisons']
  for output,dm,pm in [('lower','lower','upper'),('upper','upper','lower'),('refined_lower','refined_lower','refined_upper'),('refined_upper','refined_upper','refined_lower')]:
   numerator=Fraction();denominator=Fraction()
   for cid,arms in cc.items():
    n1=sum(g['n'] for g in arms['drug']);n0=sum(g['n'] for g in arms['placebo'])
    a=sum(g[dm] for g in arms['drug']);c=sum(g[pm] for g in arms['placebo'])
    numerator+=Fraction(a*n0,n0+n1);denominator+=Fraction(c*n1,n0+n1)
    orig=next(x for x in comps if x['comparison']==cid)
    assert n1==int(orig['drug_at_risk']) and n0==int(orig['placebo_at_risk'])
   assert numerator/denominator==Fraction(summ[ident]['exact'][output])
 report.update({'constructive_extreme_witnesses':witness,'group_outcome_rows_checked':len(groups),'analyses_reaggregated':len(by),'passed':True})
 (R/'results/validation.json').write_text(json.dumps(report,indent=2)+'\n');print(report)
if __name__=='__main__':main()