"""Frozen graph inference re-scored against the full-genomic native catalogue.

All bulky derived evidence is written to encrypted/e3008. Original artifacts stay
read-only. Gzipped historical graph inputs are restored byte-for-byte and checked
by the unchanged frozen joiner's receipt hashes before use.
"""
import gzip,hashlib,json,os,shutil,subprocess,sys,time
from pathlib import Path
B=Path(__file__).resolve().parent;CV=Path('/benchmark/pangenome-reference-cv-20261014');G=Path('/benchmark/genomic-reference-arms-20261018')
def sha(p):
 h=hashlib.sha256()
 with p.open('rb') as f:
  for chunk in iter(lambda:f.read(8*1024**2),b''):h.update(chunk)
 return h.hexdigest()
def load(p):return json.loads(p.read_text())
def save(p,r):p.write_text(json.dumps(r,indent=2)+'\n')
def restored(source,destination,names):
 destination.mkdir(parents=True)
 for name in ['COMPLETE.json',*names]:
  src=source/name;dst=destination/name
  if src.exists():dst.symlink_to(src)
  else:
   assert (source/(name+'.gz')).exists()
   with gzip.open(str(src)+'.gz','rb') as f,dst.open('wb') as out:shutil.copyfileobj(f,out)
 return destination

def main():
 assert os.environ.get('SLURM_JOB_ID');plan=load(B/'PLAN.json')
 for n,h in plan['files_sha256'].items():assert sha(B/n)==h,n
 a=plan['arms'][int(os.environ['SLURM_ARRAY_TASK_ID'])];donor=a['donor'];panel=a['panel'];fold=a['fold']
 assert sha(CV/'FOLDS.json')==plan['folds_sha256'];fs=load(CV/'FOLDS.json')['folds'][fold];assert donor in fs['test'] and donor not in fs['train']
 oldplan=CV/a['original_stage']/'PLAN.json';assert sha(oldplan)==a['original_plan_sha256'];old=load(oldplan);oldarm=next(r for r in old['arms'] if r['donor']==donor and r['panel']==panel)
 pp=CV/a['typing_stage']/(donor+'.PAIRS.json');assert sha(pp)==oldarm['pairs_sha256'];pairs=load(pp)
 source=G/'runs'/donor/('genomic_'+panel);typed=load(source/'COMPLETE.json');assert typed['status']=='complete' and typed['donor']==donor and typed['fold']==fold and typed['panel']=='genomic_'+panel and typed['pairs_sha256']==sha(pp)
 assert sha(source/'calls_genotype.tsv')==typed['native_table_sha256']
 cmd=typed['command'];assert cmd[cmd.index('--preset')+1]=='hla-wgs' and cmd[cmd.index('--alleleDigitUnits')+1]=='4' and cmd[cmd.index('--alleleDelimiter')+1]==':';reference=Path(cmd[cmd.index('-f')+1]);refreceipt=load(reference.parent/'COMPLETE.json');assert reference==G/f'fold{fold}'/'linear'/('genomic_'+panel)/'reference.fa';assert sha(reference)==typed['reference_sha256']==refreceipt['output_sha256']['reference.fa']
 aliases=load(reference.parent/'aliases.json');assert sha(reference.parent/'aliases.json')==refreceipt['output_sha256']['aliases.json']
 paneldir=CV/f'fold{fold}'/'references'/panel;assert refreceipt['source_manifest_sha256']==sha(paneldir.parent/'COMPLETE.json')
 for n,h in load(paneldir.parent/'COMPLETE.json')['output_sha256'].items():assert sha(paneldir.parent/n)==h
 for m in ['r1','r2']:assert cmd[cmd.index('-'+m[-1])+1]==pairs[m] and sha(Path(pairs[m]))==pairs[m+'_sha256']
 sys.path.insert(0,str(B/'code'))
 from context_runtime_identity import identity
 from join_evidence import build
 from run_context_evidence_disk_v3 import run as rescore
 from run_context_inference import run as infer
 from t1k_evidence import audit
 from score_native import decode_slots
 from context_inference import allele_fields
 from sequence_identity import core_slot
 output=Path(plan['output_root'])/'runs'/donor/panel
 if (output/'REBASE_COMPLETE.json').exists():
  done=load(output/'REBASE_COMPLETE.json');assert done['plan_sha256']==sha(B/'PLAN.json');print('already complete',donor,panel);return
 output.mkdir(parents=True,exist_ok=False);native=output/'native';native.mkdir()
 tool=Path(plan['export_tool']);assert sha(tool/'COMPLETE.json')==plan['export_build_sha256'];tb=load(tool/'COMPLETE.json');exe=tool/'export/genotyper';assert sha(exe)==tb['output_sha256']['export/genotyper']
 candidates=[source/f'calls_candidate_{i}.fq' for i in [1,2]];hashes={str(p):sha(p) for p in candidates}
 command=[str(exe),'-s','0.97','--alleleDigitUnits','4','--alleleDelimiter',':','--outputReadAssignment','-f',str(reference),'-1',str(candidates[0]),'-2',str(candidates[1]),'-t',os.environ['SLURM_CPUS_PER_TASK'],'-o',str(native/donor)]
 with (native/'run.log').open('wb') as log:subprocess.run(command,stdout=log,stderr=log,check=True)
 genotype=native/(donor+'_genotype.tsv');assignment=native/(donor+'_assign.tsv');assert genotype.read_bytes()==(source/'calls_genotype.tsv').read_bytes(),'Evidence export changed genotype'
 audit(assignment,native/'evidence-audit.json')
 for n,h in hashes.items():assert sha(Path(n))==h
 save(native/'COMPLETE.json',dict(status='complete',donor=donor,panel=panel,fold=fold,job_id=os.environ['SLURM_JOB_ID'],plan_sha256=sha(B/'PLAN.json'),reference_sha256=sha(reference),genotype_unchanged=True,baseline_completion_sha256=sha(source/'COMPLETE.json'),candidate_reads_sha256=hashes,input_sha256={pairs[m]:pairs[m+'_sha256'] for m in ['r1','r2']},output_sha256={p.name:sha(p) for p in [genotype,assignment,native/'evidence-audit.json']},truth_loaded=False))
 mapping=CV/a['mapping_stage']/donor;assert sha(mapping/'PLAN_TEMPLATE.json')==oldarm['mapping_template_sha256'];actual=load(mapping/'PLAN.json');template=load(mapping/'PLAN_TEMPLATE.json');assert {k:v for k,v in actual.items() if k not in ['bundle_qc_sha256','pairs_sha256']}=={k:v for k,v in template.items() if k not in ['bundle_qc_sha256','pairs_sha256']}
 if 'mapping_plan_sha256' in oldarm:assert sha(mapping/'PLAN.json')==oldarm['mapping_plan_sha256']
 roots=[]
 for gene in plan['genes']:
  root=mapping/'runs'/pairs['sample_key']/panel/gene;r=load(root/'COMPLETE.json');assert r['status']=='complete' and r['plan_sha256']==sha(mapping/'PLAN.json') and r['sample_key']==pairs['sample_key']
  dest=output/'restored_graph_inputs'/gene
  for rel,names in [('projection/fragments',['fragments.jsonl']),('projection/support',['support.jsonl']),('mapping',['alignments.jsonl'])]:restored(root/rel,dest/rel,names)
  roots.append(dest)
 calibration=CV/a['calibration_stage']/donor;cm=load(calibration/'COMPLETE.json');assert cm['status']=='calibration_complete' and cm['plan_sha256']==old['calibration_plan_sha256'];assert sha(CV/a['calibration_stage']/'PLAN.json')==old['calibration_plan_sha256'];assert sha(calibration/'calibration/COMPLETE.json')==cm['calibration_manifest_sha256']
 runtime=identity();save(output/'RUNTIME.json',runtime)
 build(native,assignment,[r/'projection/fragments' for r in roots],[r/'projection/support' for r in roots],[r/'mapping' for r in roots],output/'joined')
 rescore(output/'joined',reference,paneldir,calibration/'calibration/COMPLETE.json',[Path(pairs[m]) for m in ['r1','r2']],output/'evidence',expected_runtime=runtime)
 baseline={};native_cores={}
 for line in genotype.read_text().splitlines():
  if not line or line.startswith('#'):continue
  gene=line.split('\t')[0].removeprefix('HLA-');slots=decode_slots(line,aliases);baseline[gene]=dict(resolutions={str(f):[dict(alleles=sorted({allele_fields(x,f) for x in s['labels']}),unresolved=s['label_unknown']) for s in slots] for f in [2,4]});native_cores[gene]=[{k:s[k] for k in ['cores','core_unknown']} for s in slots]
 save(output/'baseline_calls.json',baseline);meta=load(output/'evidence/candidates.json')
 for mode in ['E1','E2','E3']:
  target=output/mode;infer(output/'evidence',output/'baseline_calls.json',target,mode);diagnostic=load(target/'decisions.json');sequence={}
  for gene in plan['genes']:
   d=diagnostic['decisions'].get(gene,{});selected=d.get('reason')=='sequence_pair_selected';slots=[core_slot(paths,meta) for paths in d['selected_paths']] if selected else native_cores.get(gene,[dict(cores=[],core_unknown=True)]*2);sequence[gene]=dict(slots=slots,origin='graph_common_score' if selected else 'native_fallback')
  save(target/'sequence_calls.json',sequence)
 save(output/'COMPLETE.json',dict(status='CV_inference_complete',donor=donor,panel=panel,fold=fold,plan_sha256=sha(B/'PLAN.json'),job_id=os.environ['SLURM_JOB_ID'],evidence_manifest_sha256=sha(output/'evidence/COMPLETE.json'),genomic_typed_receipt_sha256=sha(source/'COMPLETE.json'),output_sha256={str(p.relative_to(output)):sha(p) for mode in ['E1','E2','E3'] for p in [output/mode/'COMPLETE.json',output/mode/'calls.json',output/mode/'decisions.json',output/mode/'sequence_calls.json']},truth_loaded=False))
 subprocess.run([sys.executable,str(B/'replay.py'),'--source',str(output),'--code',str(B/'code'),'--native-table',str(genotype),'--aliases',str(reference.parent/'aliases.json'),'--output',str(output/'native_proposals_only')],check=True)
 save(output/'REBASE_COMPLETE.json',dict(status='genomic_reference_rebase_and_retrieval_ablation_complete',plan_sha256=sha(B/'PLAN.json'),genomic_inference_sha256=sha(output/'COMPLETE.json'),native_ablation_sha256=sha(output/'native_proposals_only/COMPLETE.json'),original_inference_plan_sha256=sha(oldplan),truth_loaded=False,accuracy_scored=False))
if __name__=='__main__':main()