#!/usr/bin/env python

#-------------------------------------------------------------------------

# merge_post_vcfs.py

# Merge resolved variants and other predicted variants and output the final NeuSomatic .vcf

#-------------------------------------------------------------------------

import argparse

import traceback

import fileinput

import logging

import numpy as np

from utils import get_chromosomes_order

def merge_post_vcfs(ref, resolved_vcf, no_resolve_vcf, out_vcf,

                    pass_threshold, lowqual_threshold):

    logger = logging.getLogger(merge_post_vcfs.__name__)

    logger.info("------------------------Merge vcfs-------------------------")

    chroms_order = get_chromosomes_order(reference=ref)

    good_records = []

    # Use fileinput to stream as if the two files were concatenated

    for line in fileinput.input([no_resolve_vcf, resolved_vcf]):

        if line[0] == "#":

            continue

        chrom, pos, _, ref, alt, score, _, info, format_, gt = line.strip().split()

        good_records.append([chrom, pos, ref, alt, gt, score])

    with open(out_vcf, "w") as o_f:

        o_f.write("##fileformat=VCFv4.2\n")

        o_f.write("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAMPLE\n")

        for record in sorted(good_records, key=lambda x: [chroms_order[x[0]], x[1]]):

            chrom, pos, ref, alt, gt, score = record

            prob = np.round(1 - (10**(-float(score) / 10)), 4)

            filter_ = "REJECT"

            if prob >= pass_threshold:

                filter_ = "PASS"

            elif prob >= lowqual_threshold:

                filter_ = "LowQual"

            o_f.write("\t".join([chrom, pos, ".", ref, alt,

                                 "{:.4f}".format(

                                     float(score)), filter_, "SCORE={:.4f}".format(prob),

                                 "GT", "0/1"]) + "\n")

if __name__ == '__main__':

    FORMAT = '%(levelname)s %(asctime)-15s %(name)-20s %(message)s'

    logging.basicConfig(level=logging.INFO, format=FORMAT)

    logger = logging.getLogger(__name__)

    parser = argparse.ArgumentParser(description='merge pred vcfs')

    parser.add_argument(

        '--ref', help='reference fasta filename', required=True)

    parser.add_argument('--resolved_vcf', help='resolved_vcf', required=True)

    parser.add_argument('--no_resolve_vcf',

                        help='no resolve vcf', required=True)

    parser.add_argument('--out_vcf', help='output vcf', required=True)

    args = parser.parse_args()

    logger.info(args)

    try:

        merge_post_vcfs(args.ref, args.resolved_vcf,

                        args.no_resolve_vcf, args.out_vcf)

    except Exception as e:

        logger.error(traceback.format_exc())

        logger.error("Aborting!")

        logger.error(

            "merge_post_vcfs.py failure on arguments: {}".format(args))

        raise e