#!/usr/bin/env python3
"""Convert ChromBPNet per-base contribution bigWig to BGZF/tabix BED6.

Columns: chrom, start, end, reference base, absolute contribution, sign.
Sign encodes contribution polarity, NOT genomic strand. Coordinates are
0-based, half-open. Missing/zero/nonfinite scores and non-ACGT bases are omitted.
Supports pyBigWig, or pybigtools as a fallback; pysam is required for indexing.
Use --fasta for local reference sequence, otherwise UCSC REST (hg38 by default).
"""
from __future__ import annotations
import argparse
import gzip
import json
import math
import urllib.request
from pathlib import Path

ENCODE_COUNTS_BW = 'https://encode-public.s3.amazonaws.com/2024/09/11/47f8eb25-551c-4687-ab83-158a374baa23/ENCFF829DSC.bigWig'
DYNSEQ_TRACK_IDS = ('d7a0e1c4-5f3b-4b6e-9c51-0a6f2c1d9e01','d7a0e1c4-5f3b-4b6e-9c51-0a6f2c1d9e02')

def parse_region(text):
    chrom,span=text.replace(',','').split(':');start,end=map(int,span.split('-'))
    if start<0 or end<=start:raise ValueError(f'Invalid region: {text}')
    return chrom,start,end

def merge(regions):
    out=[]
    for c,s,e in sorted(regions):
        if s<0 or e<=s:raise ValueError((c,s,e))
        if out and out[-1][0]==c and s<=out[-1][2]:out[-1]=(c,out[-1][1],max(out[-1][2],e))
        else:out.append((c,s,e))
    return out

def read_regions_bed(path,chrom_filter):
    opener=gzip.open if path.endswith('.gz') else open
    with opener(path,'rt') as f:
        regions=[]
        for line in f:
            if not line.strip() or line.startswith(('#','track','browser')):continue
            c,s,e=line.split()[:3]
            if not chrom_filter or c==chrom_filter:regions.append((c,int(s),int(e)))
    return merge(regions)

class Sequence:
    def __init__(self,fasta,genome,cache=None):
        self.genome=genome;self.fa=None;self.cache=json.loads(Path(cache).read_text()) if cache else None
        if fasta:
            import pyfaidx
            self.fa=pyfaidx.Fasta(fasta,as_raw=True,sequence_always_upper=True)
    def fetch(self,chrom,start,end):
        if self.fa is not None:return str(self.fa[chrom][start:end]).upper()
        if self.cache is not None:
            c=self.cache
            assert c['genome']==self.genome and c['chrom']==chrom and c['start']<=start<end<=c['end']
            return c['dna'][start-c['start']:end-c['start']].upper()
        url=f'https://api.genome.ucsc.edu/getData/sequence?genome={self.genome};chrom={chrom};start={start};end={end}'
        with urllib.request.urlopen(url,timeout=60) as r:return json.load(r)['dna'].upper()

def open_bigwig(path):
    try:
        import pyBigWig
        if not path.startswith(('http://','https://')) or pyBigWig.remote:
            return pyBigWig.open(path),'pyBigWig'
    except ImportError:pass
    import pybigtools
    return pybigtools.open(path),'pybigtools'

def patch_domains(spec,ymax):
    count=0
    def visit(node):
        nonlocal count
        if isinstance(node,dict):
            if node.get('id') in DYNSEQ_TRACK_IDS and isinstance(node.get('y'),dict):
                node['y']['domain']=[0,ymax];count+=1
            for v in node.values():visit(v)
        elif isinstance(node,list):
            for v in node:visit(v)
    visit(spec)
    if count!=2:raise ValueError(f'Expected 2 dynseq track IDs, found {count}')
    return spec

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--bigwig',default=ENCODE_COUNTS_BW)
    w=p.add_mutually_exclusive_group(required=True);w.add_argument('--region',action='append');w.add_argument('--regions-bed')
    p.add_argument('--chrom');p.add_argument('--fasta');p.add_argument('--genome',default='hg38')
    p.add_argument('--sequence-json',help='Optional cached UCSC sequence JSON for reproducible/offline reference lookup')
    p.add_argument('--out',required=True);p.add_argument('--min-abs',type=float,default=0)
    p.add_argument('--domain-quantile',type=float,default=1.0)
    p.add_argument('--spec-in');p.add_argument('--spec-out');args=p.parse_args()
    if not math.isfinite(args.min_abs) or args.min_abs<0:p.error('--min-abs must be finite and nonnegative')
    if not 0<args.domain_quantile<=1:p.error('--domain-quantile must be in (0,1]')
    if args.spec_out and not args.spec_in:p.error('--spec-out requires --spec-in')
    regions=merge(parse_region(r) for r in args.region) if args.region else read_regions_bed(args.regions_bed,args.chrom)
    if not regions:p.error('No regions selected')
    bed=Path(args.out[:-3] if args.out.endswith('.gz') else args.out)
    for path in [bed,Path(str(bed)+'.gz'),Path(str(bed)+'.gz.tbi')]:
        if path.exists():p.error(f'Refusing to overwrite {path}')
    bed.parent.mkdir(parents=True,exist_ok=True)
    bw,backend=open_bigwig(args.bigwig);chroms=bw.chroms();seq=Sequence(args.fasta,args.genome,args.sequence_json)
    mags=[];positive=0;covered=0
    try:
        with bed.open('w') as out:
            for chrom,start,end in regions:
                if chrom not in chroms or end>chroms[chrom]:raise ValueError(f'Region outside bigWig chromosome bounds: {chrom}:{start}-{end}')
                for a in range(start,end,100000):
                    b=min(end,a+100000)
                    vals=bw.values(chrom,a,b,numpy=False) if backend=='pyBigWig' else bw.values(chrom,a,b,missing=float('nan'))
                    bases=seq.fetch(chrom,a,b)
                    assert len(vals)==len(bases)==b-a
                    for i,(v,base) in enumerate(zip(vals,bases)):
                        if v is None or not math.isfinite(v):continue
                        covered+=1
                        if abs(v)<=args.min_abs or base not in 'ACGT':continue
                        sign='+' if v>0 else '-';mag=abs(float(v))
                        out.write(f'{chrom}\t{a+i}\t{a+i+1}\t{base}\t{mag:.9g}\t{sign}\n')
                        mags.append(mag);positive+=v>0
    finally:bw.close()
    if not mags:
        bed.unlink();raise ValueError('No nonzero A/C/G/T contribution scores in selected regions')
    mags.sort();k=min(len(mags)-1,math.ceil(args.domain_quantile*len(mags))-1)
    # Round UP so the full-range domain cannot clip a maximum due to rounding.
    value=mags[k];step=10**(math.floor(math.log10(value))-3);ymax=math.ceil(value/step)*step
    import pysam
    gz=pysam.tabix_index(str(bed),preset='bed',force=False,keep_original=False)
    with pysam.TabixFile(gz) as t:assert sum(1 for _ in t.fetch())==len(mags)
    report=dict(bigwig=args.bigwig,assembly=args.genome,regions=regions,rows=len(mags),positive=int(positive),negative=len(mags)-int(positive),finite_scored_bases=covered,max_abs_score=mags[-1],domain_quantile=args.domain_quantile,y_domain=[0,ymax],backend=backend,sign_is_contribution_not_strand=True)
    Path(str(bed)+'.summary.json').write_text(json.dumps(report,indent=2)+'\n')
    if args.spec_in:
        spec=patch_domains(json.loads(Path(args.spec_in).read_text()),ymax)
        Path(args.spec_out or args.spec_in).write_text(json.dumps(spec,indent=2)+'\n')
    print(json.dumps(report,indent=2));print('Wrote',gz,'and index')

if __name__=='__main__':main()