#!/usr/bin/env python

# modules
import argparse
from time import asctime as at
import sys
import itertools

printHelp = "off"
if len(sys.argv) < 2:
	printHelp = "on"
elif len(sys.argv) >= 2:
	if (sys.argv[1] in ["-h", "--help"]):
		printHelp = "on"

if printHelp == "on":
	sys.exit('''

SNP-Dogg

last update: 20/09/2017
---------------------------------------------------------------------------------------------------------------------------------

-in	--input-file		VCF Format (not BCF) [stdin]
-v	--verbose		Comment steps in the stderr [off]
-q	--quality		Minimum call quality (VCF field 6) [off]
-amq	--avg-map-qual		Minimum average mapping quality [off]
-d	--min-depth		Minimum coverage of the SNP (based on DP4 field) [off]
-D	--max-depth		... and maximum [100000]
-mqzf	--map-qual-zero-frac	Maximum tolerated fraction of reads with mapping quality 0 (float representing %%) [off]
-af	--alt-frac		% of reads confirming alternative allele [off]
-mr	--min-read		Minimum number of reads confirming alternative allele [off]
-bs	--both-strands		Only together with \"--alt-frac\" or \"--min-read\". Require same criterion on both strands separately [off]


''')

# parser
p = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, add_help=False)
p.add_argument("-in", "--input-file")
p.add_argument("-v", "--verbose", action="store_true")
p.add_argument("-q", "--quality", default = 0, type=float)
p.add_argument("-amq", "--avg-map-qual", type=int, default=0)
p.add_argument("-d", "--min-depth", default = 0, type=int)
p.add_argument("-D", "--max-depth", default = 100000, type=int)
p.add_argument("-mqzf", "--map-qual-zero-frac", default=100, type=float)
p.add_argument("-af", "--alt-frac", default=0, type=float)
p.add_argument("-bs", "--both-strands", action="store_true")
p.add_argument("-mr", "--min-read", default=0, type=int)
args = p.parse_args()

# verifications
if args.both_strands and not (args.alt_frac or args.min_read):
	sys.exit("[{0}]\tERROR: you can specify --both-strands only if you use --alt-freq or --min-read as well.\n".format(at()))

# read input VCF file, print filtered header, do filtering 
if args.verbose:
	sys.stderr.write("[{0}]\tFiltering VCF file\n".format(at()))

if args.input_file:
	INPUT = open(args.input_file, "r")
else:
	INPUT = sys.stdin

j=0
for line in INPUT:
	if line[0:1] == "#":
		sys.stdout.write("{0}\n".format(line.rstrip()))
	else:
		# filter lines according to criteria
		printSwitch = "off"
		# declare dictionary of the info field
		INFO = {}
		lst = line.rstrip().split("\t")
		temp = lst[7].split(";")
		if temp[0] == "INDEL":
			del temp[0]
		for item in temp:
			INFO[item.split("=")[0]] = item.split("=")[1]
		# do the selection on the DP4 field
		temp = INFO["DP4"].split(",")
		ref = float(temp[0]) + float(temp[1])
		alt = float(temp[2]) + float(temp[3])
		INFO["DP"] = ref + alt
		ref_f = float(temp[0])
		ref_r = float(temp[1])
		alt_f = float(temp[2])
		alt_r = float(temp[3])
		try:
			ratio = float(alt/(alt+ref))*100
		except ZeroDivisionError:
			ratio = 0
		try:
			ratio_f = float(alt_f/(ref_f+alt_f))*100
		except ZeroDivisionError:
			ratio_f = 0
		try:
			ratio_r = float(alt_r/(ref_r+alt_r))*100
		except ZeroDivisionError:
			ratio_r = 0
		if (args.min_depth <= int(INFO["DP"]) <= args.max_depth):
			if (float(lst[5]) >= args.quality):
				if (float(INFO["MQ0F"])*100 <= args.map_qual_zero_frac):
					if (float(INFO["MQ"]) >= args.avg_map_qual):
						if args.both_strands:
							if ((ratio_r >= args.alt_frac) and (ratio_f >= args.alt_frac)) ^ ((alt_r >= args.min_read) and (alt_f >= args.min_read)):
								printSwitch = "on"
						else:
							if ((ratio >= args.alt_frac) and (alt >= args.min_read)):
								printSwitch = "on"
		if printSwitch == "on":
			j+=1
			sys.stdout.write(line) 
#
INPUT.close()
if args.verbose:
	sys.stderr.write("[{0}]\tSelected {1} variants\n".format(at(), j))
