#!/usr/bin/env python3
"""Run pinned ART-modern, synchronize its unsorted mates, and gzip them."""
import argparse,pathlib,subprocess,gzip,hashlib,json,collections,time,os
P=pathlib.Path(__file__).resolve().parent
ap=argparse.ArgumentParser();ap.add_argument('--input',required=True);ap.add_argument('--output-prefix',required=True);ap.add_argument('--seed',type=int,required=True);ap.add_argument('--threads',type=int,default=8);ap.add_argument('--coverage',type=float,default=15);ap.add_argument('--name-prefix',required=True);a=ap.parse_args()
assert 1<=a.seed<=2147483647 and a.threads>=1 and a.coverage>0
assert all(c.isalnum() or c in '_.-' for c in a.name_prefix)
out=pathlib.Path(a.output_prefix);out.parent.mkdir(parents=True,exist_ok=True)
paths=[str(out)+f'_R{i}.fastq.gz' for i in (1,2)];receipt=str(out)+'.COMPLETE.json'
assert not any(pathlib.Path(x).exists() for x in paths+[receipt]),'refuse overwrite'
cmd=[str(P/'bin-1.3.2/art_modern'),'--mode','wgs','--lc','pe','--i-file',str(pathlib.Path(a.input).resolve()),'--i-fcov',str(a.coverage),'--read_len','150','--pe_frag_dist_mean','350','--pe_frag_dist_std_dev','50','--i-seed',str(a.seed),'--parallel',str(a.threads),'--qual_file_1',str(P/'NovaSeqX_ERR15382247_R1.txt'),'--qual_file_2',str(P/'NovaSeqX_ERR15382247_R2.txt'),'--ins_rate_1','0','--ins_rate_2','0','--del_rate_1','0','--del_rate_2','0','--max_indel','0','--max_n','0','--min_qual','0','--max_qual','71','--o-fastq','/dev/stdout']
env=os.environ.copy();env['ART_LOG_DIR']=str(out)+'.logs';pending={};n=0;maxpending=0;hist=[collections.Counter(),collections.Counter()];start=time.time()
with open(str(out)+'.stderr','wb') as err,gzip.open(paths[0],'wb',compresslevel=1) as o1,gzip.open(paths[1],'wb',compresslevel=1) as o2:
 proc=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=err,env=env)
 try:
  while True:
   h=proc.stdout.readline()
   if not h:break
   s,plus,q=[proc.stdout.readline() for _ in range(3)]
   assert h.startswith(b'@') and plus.startswith(b'+') and len(s.rstrip())==len(q.rstrip())==150
   key,mate=h.rstrip().rsplit(b'/',1);mi=int(mate)-1;assert mi in (0,1)
   rec=b'@'+a.name_prefix.encode()+b':'+key[1:]+b'/'+mate+b'\n'+s+plus+q
   hist[mi].update(c-33 for c in q.rstrip())
   if key in pending:
    prevmi,prev=pending.pop(key);assert prevmi!=mi
    pair=[None,None];pair[mi]=rec;pair[prevmi]=prev;o1.write(pair[0]);o2.write(pair[1]);n+=1
   else:pending[key]=(mi,rec)
   maxpending=max(maxpending,len(pending));assert maxpending<=1000000,'unexpected unsynchronized output'
  assert proc.wait()==0,'ART failed; partial outputs retained'
  assert not pending and n>0
 except BaseException:
  proc.kill();proc.wait();raise

def sha(path):
 h=hashlib.sha256()
 with open(path,'rb') as f:
  for b in iter(lambda:f.read(1048576),b''):h.update(b)
 return h.hexdigest()
r=dict(status='complete',command=cmd,input_sha256=sha(a.input),binary_sha256=sha(P/'bin-1.3.2/art_modern'),profiles_sha256=[sha(P/f'NovaSeqX_ERR15382247_R{i}.txt') for i in (1,2)],wrapper_sha256=sha(__file__),pairs=n,bases=n*300,output_paths=paths,output_sha256=[sha(x) for x in paths],quality_histograms=[dict(x) for x in hist],maximum_pending_mates=maxpending,seconds=time.time()-start,model='empirical NovaSeq X Plus prefix-cycle mate quality marginals; Phred-derived substitutions; zero indels assumption')
pathlib.Path(receipt).write_text(json.dumps(r,indent=2)+'\n')