"""Verify linkage inference against the frozen likelihood and prior receipts."""
import hashlib,json,subprocess
from pathlib import Path
from model import candidate_prior
B=Path(__file__).resolve().parent
sha=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
load=lambda p:json.loads(p.read_text())
SSH=['ssh','-F','/runtime/user/.ssh/config','-o','BatchMode=yes','-o','ConnectTimeout=10','-o','LogLevel=ERROR']
def main():
 p=load(B/'PLAN.json');s=load(B/'SUBMISSION.json');assert sha(B/'PLAN.json')==s['plan_sha256']
 for name,h in p['inputs_sha256'].items():assert sha(B/name)==h
 lines=subprocess.check_output(SSH+['ibex','sacct -X -n -P -j '+s['job_id']+' --format=JobID,JobIDRaw,State,ExitCode'],text=True).splitlines();states={x.split('|')[0]:x.split('|')[1:] for x in lines};ready=[]
 for i,arm in enumerate(p['arms']):
  state=states.get(s['job_id']+'_'+str(i))
  if state and state[1:3]==['COMPLETED','0:0']:ready.append((arm,state[0]))
 listing=B/'OUTPUT_FILES.txt';listing.write_text(''.join(a['donor']+'/'+a['panel']+'/\n' for a,_ in ready));dest=B/'collected';dest.mkdir(exist_ok=True)
 if ready:subprocess.run(['rsync','-ar','--files-from='+str(listing),'-e',' '.join(SSH),'ibex:'+s['remote']+'/runs/',str(dest)+'/'],check=True)
 lb=B.parent/p['likelihood_root'];pb=B.parent/p['prior_root'];lv=load(lb/'COLLECTION.json');pv=load(pb/'COLLECTION.json');verified=[];awaiting=[]
 for arm,job in ready:
  key=(arm['donor'],arm['panel']);source=next((x for x in lv['verified'] if (x['donor'],x['panel'])==key),None);priorproof=next((x for x in pv['verified'] if x['panel']['file']==arm['training_panel']),None)
  if source is None or priorproof is None:awaiting.append(arm);continue
  d=dest/arm['donor']/arm['panel'];r=load(d/'RESULT.json');lr=lb/'collected'/arm['donor']/arm['panel'];pr=pb/'collected'/(arm['training_panel']+'.RESULT.json');prior=load(pr);panel=load(lb/arm['training_panel'])
  assert r['status']=='shared_read_linkage_ablation_inference_complete' and r['arm']==arm and r['job_id']==job and r['plan_sha256']==s['plan_sha256']
  assert r['likelihood_receipt_sha256']==source['result_sha256']==sha(lr/'RESULT.json')
  assert r['prior_selection_sha256']==priorproof['result_sha256']==sha(pr)
  assert not r['truth_loaded'] and not r['accuracy_scored']
  q=[]
  for i,g in enumerate(panel['genes']):
   record=load(lr/(g+'.json'));assert sha(lr/(g+'.json'))==load(lr/'RESULT.json')['output_sha256'][g+'.json']
   q.append(candidate_prior([x['identity'] for x in record['candidates']],[x['core_ids'][i] for x in panel['observed_associations']],prior['alpha']))
  for mode in ['independent','linked','shuffled']:
   name=mode+'.json';assert sha(d/name)==r['output_sha256'][name];result=load(d/name)
   assert result['mode']==mode and result['candidate_prior']==q and result['unlinked_mass']==prior['selected']['unlinked_mass']
   assert result['rho']==(1 if mode=='independent' else prior['selected']['rho'])
   assert not result['truth_loaded'] and not result['phase_decoded'] and not result['posterior_calibrated']
   for g in panel['genes']:
    call=result['calls'][g];record=load(lr/(g+'.json'));n=len(record['candidates'])
    assert call['assigned_fragments']==record['fragments'] and call['imputed_without_assigned_reads']==(record['fragments']==0)
    assert 0<=call['unphased_pair_score']<=1+1e-12 and len(call['slots'])==2
    assert call['tied_pairs']==len(call['best_candidate_pairs'])>0
    assert all(0<=i<=j<n for i,j in call['best_candidate_pairs'])
  verified.append({'arm':arm,'result_sha256':sha(d/'RESULT.json')})
 report={'verified':verified,'planned':len(p['arms']),'awaiting_source_verification':awaiting,'scheduler':lines,'accuracy_scored':False};(B/'COLLECTION.json').write_text(json.dumps(report,indent=2)+'\n');print(json.dumps({'verified':len(verified),'awaiting_sources':len(awaiting),'planned':len(p['arms'])}))
if __name__=='__main__':main()