#!/usr/bin/env python3
"""Convert all bigWig-scored hg38 bases to BED6/BGZF, with resumable chromosome jobs."""
import argparse,json,math,time,subprocess,shutil,concurrent.futures,itertools
from pathlib import Path
import pysam,pybigtools

def convert(job):
 source,fasta,chrom,length,parts=job; path=Path(parts)/(chrom+'.bgz'); meta=Path(str(path)+'.json')
 if path.exists() and meta.exists(): return chrom,json.loads(meta.read_text())
 bw=pybigtools.open(source); fa=pysam.FastaFile(fasta); records=bw.records(chrom);first=next(records,None)
 stats={'rows':0,'positive':0,'negative':0,'max_absolute':0.,'max_positive':0.}
 if first is None: meta.write_text(json.dumps(stats));return chrom,stats
 if chrom not in fa.references or fa.get_reference_length(chrom)!=length: raise ValueError('Reference mismatch: '+chrom)
 seq=fa.fetch(chrom).upper(); count=pos=0; mx=pmx=0.; buf=[]
 with open(path,'wb') as compressed:
  compressor=subprocess.Popen(['bgzip','-c','-l','1'],stdin=subprocess.PIPE,stdout=compressed)
  try:
   for start,end,value in itertools.chain([first],records):
    if not math.isfinite(value) or value==0: continue
    mag=abs(value); positive=value>0; sign='+' if positive else '-'; vs=format(mag,'.9g')
    mx=max(mx,mag)
    if positive: pmx=max(pmx,mag)
    for i in range(start,end):
     base=seq[i]
     if base not in 'ACGT': continue
     buf.append(f'{chrom}\t{i}\t{i+1}\t{base}\t{vs}\t{sign}\n');count+=1;pos+=positive
     if len(buf)>=20000: compressor.stdin.write(''.join(buf).encode());buf.clear()
   if buf: compressor.stdin.write(''.join(buf).encode())
  finally: compressor.stdin.close()
  if compressor.wait()!=0: raise RuntimeError('bgzip failed')
 stats.update(rows=count,positive=pos,negative=count-pos,max_absolute=mx,max_positive=pmx)
 meta.write_text(json.dumps(stats));return chrom,stats

def main():
 p=argparse.ArgumentParser();p.add_argument('bigwig');p.add_argument('out');p.add_argument('--workers',type=int,default=8);p.add_argument('--fasta',default='/net/users/voy/nmf/data_itcr/cnTFeffects/script1_output/hg38.fa');a=p.parse_args()
 chroms=pybigtools.open(a.bigwig).chroms();parts=Path(a.out+'.parts');parts.mkdir(exist_ok=True)
 report={'input':a.bigwig,'fasta':a.fasta,'chromosomes':{},'coordinate_system':'0-based half-open','columns':['chrom','start','end','reference_base','absolute_contribution','contribution_sign']};t0=time.time()
 jobs=[(a.bigwig,a.fasta,c,n,str(parts)) for c,n in sorted(chroms.items(),key=lambda x:-x[1])]
 with concurrent.futures.ProcessPoolExecutor(a.workers) as pool:
  futures=[pool.submit(convert,j) for j in jobs]
  for f in concurrent.futures.as_completed(futures):
   c,s=f.result();report['chromosomes'][c]=s
   if s['rows']: print(c,s['rows'],'bases',round(time.time()-t0),'seconds',flush=True)
 # BGZF EOF blocks must only occur at the very end of the concatenation.
 eof=bytes.fromhex('1f8b08040000000000ff0600424302001b0003000000000000000000')
 with open(a.out,'wb') as dest:
  for chrom in sorted(chroms):
   if not report['chromosomes'][chrom]['rows']:continue
   part=parts/(chrom+'.bgz')
   with part.open('rb') as src:
    src.seek(-28,2);assert src.read()==eof
    remaining=part.stat().st_size-28;src.seek(0)
    while remaining:
     data=src.read(min(8*1024*1024,remaining));dest.write(data);remaining-=len(data)
  dest.write(eof)
 print('Indexing',flush=True);pysam.tabix_index(a.out,preset='bed',force=True)
 print('Validating every chromosome row count',flush=True)
 with pysam.TabixFile(a.out) as tb:
  for chrom,stats in report['chromosomes'].items():
   n=sum(1 for _ in tb.fetch(chrom)) if chrom in tb.contigs else 0
   assert n==stats['rows'],(chrom,n,stats)
 report['validated']=True;report['elapsed_seconds']=time.time()-t0
 Path(a.out+'.summary.json').write_text(json.dumps(report,indent=2)+'\n');print('COMPLETE',a.out,flush=True)
if __name__=='__main__':main()