"""Compare independent, observed-linked and shuffled priors on identical reads."""
import collections,hashlib,json,math,os
from pathlib import Path
from model import candidate_prior,decode
from open_linkage import infer
B=Path(__file__).resolve().parent;ROOT=Path('/benchmark')
sha=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
load=lambda p:json.loads(p.read_text())
def main():
 assert os.environ.get('SLURM_JOB_ID');plan=load(B/'PLAN.json')
 for n,h in plan['inputs_sha256'].items():assert sha(B/n)==h
 arm=plan['arms'][int(os.environ['SLURM_ARRAY_TASK_ID'])];lb=ROOT/plan['likelihood_root'];pb=ROOT/plan['prior_root']
 assert sha(lb/'PLAN.json')==plan['likelihood_plan_sha256'] and sha(pb/'PLAN.json')==plan['prior_plan_sha256']
 panel=load(lb/arm['training_panel']);pr=load(pb/(arm['training_panel']+'.RESULT.json'))
 assert pr['status']=='training_family_prior_selection_complete' and pr['plan_sha256']==plan['prior_plan_sha256']
 assert pr['training_panel_sha256']==sha(lb/arm['training_panel']) and not pr['outer_test_loaded']
 assert arm['donor'] in panel['excluded_donors'] and arm['donor'] not in panel['training_donors']
 winner=max(pr['candidates'],key=lambda r:(r['mean_family_log_predictive_score'],r['rho'],r['unlinked_mass']))
 assert all(winner[k]==v for k,v in pr['selected'].items())
 evidence=lb/'runs'/arm['donor']/arm['panel'];receipt=load(evidence/'RESULT.json')
 assert receipt['status']=='shared_linkage_read_likelihoods_complete' and receipt['arm']==arm and receipt['plan_sha256']==plan['likelihood_plan_sha256']
 assert receipt['training_panel_sha256']==pr['training_panel_sha256'] and not receipt['truth_loaded']
 records=[];matrices=[];priors=[];templates=[];shuffled=[];genes=panel['genes'];h=len(panel['observed_associations'])
 assert h==len(panel['shuffled_associations'])
 for i,gene in enumerate(genes):
  name=gene+'.json';assert sha(evidence/name)==receipt['output_sha256'][name];r=load(evidence/name);records.append(r)
  identities=[c['identity'] for c in r['candidates']];index={x:j for j,x in enumerate(identities)};assert len(index)==len(identities)
  col=[x['core_ids'][i] for x in panel['observed_associations']];other=[x['core_ids'][i] for x in panel['shuffled_associations']];assert collections.Counter(col)==collections.Counter(other)
  templates.append([index[x] for x in col]);shuffled.append([index[x] for x in other]);matrices.append(r['scaled_likelihood']);priors.append(candidate_prior(identities,col,pr['alpha']))
 templates=list(map(list,zip(*templates)));shuffled=list(map(list,zip(*shuffled)));weights=[1/h]*h
 out=B/'runs'/arm['donor']/arm['panel'];out.mkdir(parents=True,exist_ok=False);outputs={};parameters=pr['selected'];summary=[]
 for mode,states,rho in [('independent',templates,1),('linked',templates,parameters['rho']),('shuffled',shuffled,parameters['rho'])]:
  result=infer(matrices,states,weights,priors,parameters['unlinked_mass'],[rho]*(len(genes)-1));calls={}
  for i,gene in enumerate(genes):
   calls[gene]=dict(decode(result['ordered_allele_pair_marginals'][i],records[i]['candidates'],plan['tie_tolerance']),assigned_fragments=records[i]['fragments'],imputed_without_assigned_reads=records[i]['fragments']==0)
  payload={'mode':mode,'calls':calls,'rho':rho,'unlinked_mass':parameters['unlinked_mass'],'log_evidence_up_to_common_read_scale':result['log_evidence'],'candidate_prior':priors,'posterior_calibrated':False,'phase_decoded':False,'truth_loaded':False}
  p=out/(mode+'.json');p.write_text(json.dumps(payload,indent=2)+'\n');outputs[p.name]=sha(p);summary.append({'mode':mode,'rho':rho,'imputed_loci':sum(x['imputed_without_assigned_reads'] for x in calls.values())})
 report={'status':'shared_read_linkage_ablation_inference_complete','arm':arm,'job_id':os.environ['SLURM_JOB_ID'],'plan_sha256':sha(B/'PLAN.json'),'likelihood_receipt_sha256':sha(evidence/'RESULT.json'),'prior_selection_sha256':sha(pb/(arm['training_panel']+'.RESULT.json')),'output_sha256':outputs,'summary':summary,'truth_loaded':False,'accuracy_scored':False,'limitations':['Development pilot, conditional on uniquely assigned internal read fragments and fixed candidate retrieval.','Independent, linked and shuffled arms share likelihoods and marginal candidate priors.','Prior parameters selected with training family holdouts; read-error and posterior calibration unestablished.','Unphased marginal typing only; no chromosome phase path or within-gene block reconstruction.','Sequences absent from every candidate are not recoverable.']}
 (out/'RESULT.json').write_text(json.dumps(report,indent=2)+'\n');print(json.dumps({'status':report['status'],'summary':summary}))
if __name__=='__main__':main()