# forked from https://github.com/single-cell-genetics/cellSNP
## A python wrap of UCSC liftOver function for vcf file
## UCSC liftOver binary and hg19 to hg38 chain file:
## https://genome.ucsc.edu/cgi-bin/hgLiftOver
## http://hgdownload.cse.ucsc.edu/admin/exe/linux.x86_64/liftOver
## http://hgdownload.soe.ucsc.edu/goldenPath/hg19/liftOver/hg19ToHg38.over.chain.gz

import sys
import gzip
import subprocess
from optparse import OptionParser

LIFTOVER_INFO = '##INFO=<ID=OLD,Number=1,Type=Integer,'
LIFTOVER_INFO += 'Description="position before liftover">\n'

def vcf_to_bed(vcf_file, out_file, chr_in=True):
    if vcf_file[-3:] == ".gz":
        is_gzip = True
        fid_in = gzip.open(vcf_file, "r")
    else:
        is_gzip = False
        fid_in = open(vcf_file, "r")
    
    fid_out = open(out_file, "w")
    for line in fid_in:
        if is_gzip:
            line = line.decode('utf-8')
        if line.startswith("#") == False:
            line_val = line.rstrip().split("\t")[:8]
            if chr_in and line_val[0].startswith("chr") == False:
                line_val[0] = "chr" + line_val[0]
            line_val[2] = str(int(line_val[1]) + 1)
            fid_out.writelines("\t".join(line_val[:3]) + "\n")
    fid_in.close()
    fid_out.close()
    return None

def update_vcf(vcf_file, bed_new, bed_unmap, out_file):       
    ## unmapped lines
    unmap_pos = []
    _fid = open(bed_unmap, "r")
    for line in _fid:
        if not line.startswith("#"):
            _pos_id = "_".join(line.rstrip().split("\t")[:2])
            unmap_pos.append(_pos_id)
    _fid.close()
    
    if vcf_file[-3:] == ".gz":
        is_gzip = True
        fid_in = gzip.open(vcf_file, "r")
    else:
        is_gzip = False
        fid_in = open(vcf_file, "r")
        
    cnt1 = 0 
    idx_unmap = 0
    fid_bed = open(bed_new, "r")
    fid_out = open(out_file, "w")    
    for line in fid_in:
        if is_gzip:
            line = line.decode('utf-8')
        if line.startswith("#"):
            if line.startswith("#CHROM"):
                fid_out.writelines(LIFTOVER_INFO)
            fid_out.writelines(line)
        else:
            line_val = line.rstrip().split("\t")
            if idx_unmap < len(unmap_pos):
                _pos_id = "_".join(line_val[:2])
                if line_val[0].startswith("chr") == False:
                    _pos_id = "chr" + _pos_id
                if _pos_id == unmap_pos[idx_unmap]:
                    idx_unmap += 1
                    continue
            cnt1 += 1
            bed_line = fid_bed.readline()
            line_val[7] = "OLD=" + line_val[1] + ";" + line_val[7]
            line_val[1] = bed_line.rstrip().split("\t")[1]
            fid_out.writelines("\t".join(line_val) + "\n")
    print(cnt1, idx_unmap)
    fid_in.close()
    fid_bed.close()
    fid_out.close()
    return None


def main():
    import warnings
    warnings.filterwarnings('error')

    # parse command line options
    parser = OptionParser()
    parser.add_option("--chainFile", "-c", dest="chain_file", default=None,
        help=("Chain file, full path."))
    parser.add_option("--inFile", "-i", dest="in_file", default=None,
        help=("Input vcf file, full path."))
    parser.add_option("--outFile", "-o", dest="out_file", default=None,
        help=("Output VCF file, full path."))
    parser.add_option("--liftOverPath", "-P", dest="liftOver_path", default=None,
        help=("liftOver_path if it is not in PATH variable."))
    
    (options, args) = parser.parse_args()
    if len(sys.argv[1:]) == 0:
        print("liftOver-vcf: a wrap of UCSC liftOver for VCF file.\n")
        print("use -h or --help for help on argument.")
        sys.exit(1)
        
    in_file = options.in_file
    bed_file = options.in_file.split(".vcf")[0] + ".bed"
    new_bed_file = options.out_file.split(".vcf")[0] + ".bed"
    unmap_bed_file = options.out_file.split(".vcf")[0] + ".unmap.bed"
    
    ## generate bed file
    print("converting vcf to bed file ... ")
    vcf_to_bed(in_file, bed_file)
    
    ## UCSC liftOver on bed file
    chain_file = options.chain_file
    if options.liftOver_path is None:
        liftOver = "liftOver"
    else:
        # check if path exists
        liftOver = options.liftOver_path
                
    print("liftOver bed file ... ")
    bashCommand = "%s %s %s %s %s" %(liftOver, bed_file, chain_file, 
                                  new_bed_file, unmap_bed_file)
    #print(bashCommand)
    pro = subprocess.Popen(bashCommand.split(), stdout=subprocess.PIPE)
    pro.communicate()[0]

    ## update vcf file
    out_file = options.out_file
    if out_file[-3:] == ".gz":
        out_file = out_file[:-3]
    print("updating vcf file ... ")
    update_vcf(in_file, new_bed_file, unmap_bed_file, out_file)
    
    print("gzip vcf file ... ")
    import shutil
    if shutil.which("bgzip") is not None:
        bashCommand = "bgzip -f %s" %(out_file)
    else:
        bashCommand = "gzip -f %s" %(out_file)
    pro = subprocess.Popen(bashCommand.split(), stdout=subprocess.PIPE)
    pro.communicate()[0]
    return None

if __name__ == "__main__":
    main()
    

[((89, 4, 89, 36), 'warnings.filterwarnings', 'warnings.filterwarnings', ({(89, 28, 89, 35): '"""error"""'}, {}), "('error')", False, 'import warnings\n'), ((92, 13, 92, 27), 'optparse.OptionParser', 'OptionParser', ({}, {}), '()', False, 'from optparse import OptionParser\n'), ((19, 17, 19, 41), 'gzip.open', 'gzip.open', ({(19, 27, 19, 35): 'vcf_file', (19, 37, 19, 40): '"""r"""'}, {}), "(vcf_file, 'r')", False, 'import gzip\n'), ((50, 17, 50, 41), 'gzip.open', 'gzip.open', ({(50, 27, 50, 35): 'vcf_file', (50, 37, 50, 40): '"""r"""'}, {}), "(vcf_file, 'r')", False, 'import gzip\n'), ((106, 8, 106, 19), 'sys.exit', 'sys.exit', ({(106, 17, 106, 18): '(1)'}, {}), '(1)', False, 'import sys\n'), ((141, 7, 141, 28), 'shutil.which', 'shutil.which', ({(141, 20, 141, 27): '"""bgzip"""'}, {}), "('bgzip')", False, 'import shutil\n')]