import sys
import argparse
import logging
import time
import vcf

def call_gt(tag):
	if sum(tag) == 2:
		return '1/1'
	else:
		if tag[0] == 1:
			return '1/0'
		if tag[1] == 1:
			return '0/1'

	return './.'

def main_ctrl(args):

	fileout = open(args.outvcf, 'w')

	header = open(args.invcf, 'r')
	for line in header:
		if line[0] == '#':
			fileout.write(line)
		else:
			continue
	header.close()

	filein = vcf.Reader(open(args.invcf, 'r'))
	for record in filein:
		if len(record.FILTER) == 0:
			if record.INFO["RE"] == 1:
				try:
					fileout.write("{chr}\t{pos}\t{id}\t{ref}\t{alt}\t{qual}\t{filter}\t{info}\t{format}\t{gt}\n".format(
						chr = record.CHROM,
						pos = record.POS,
						id = record.ID,
						ref = record.REF,
						alt = record.ALT[0],
						qual = record.QUAL,
						filter = 'PASS',
						info = "SVTYPE=%s;SVLEN=%d;END=%d;RE=%d;RNAMES=%s"%(record.INFO['SVTYPE'], 
							record.INFO['SVLEN'], 
							record.INFO['END'], 
							record.INFO['RE'], 
							','.join(record.INFO['RNAMES'])),
						format = 'GT',
						gt = '1/0' if 'cutesvh1' in record.INFO['RNAMES'][0] else '0/1'
						))
				except:
					pass

			else:
				tag = [0, 0]
				for i in record.INFO['RNAMES']:
					if 'cutesvh1' in i:
						tag[0] = 1
					if 'cutesvh2' in i:
						tag[1] = 1

				try:
					fileout.write("{chr}\t{pos}\t{id}\t{ref}\t{alt}\t{qual}\t{filter}\t{info}\t{format}\t{gt}\n".format(
						chr = record.CHROM,
						pos = record.POS,
						id = record.ID,
						ref = record.REF,
						alt = record.ALT[0],
						qual = record.QUAL,
						filter = 'PASS',
						info = "SVTYPE=%s;SVLEN=%d;END=%d;RE=%d;RNAMES=%s"%(record.INFO['SVTYPE'], 
							record.INFO['SVLEN'], 
							record.INFO['END'], 
							record.INFO['RE'], 
							','.join(record.INFO['RNAMES'])),
						format = 'GT',
						# gt = '1/1' if sum(tag) == 2 else '0/1'
						gt = call_gt(tag)
						))
				except:
					pass
	fileout.close()

def main(argv):
	args = parseArgs(argv)
	setupLogging(False)
	starttime = time.time()
	main_ctrl(args)
	logging.info("Finished in %0.2f seconds."%(time.time() - starttime))

USAGE="""\
	Convert the typical SV callsets generated from cuteSV to diploid based SV callsets.
	Author: Tao Jiang
	Email: tjiang@hit.edu.cn
"""

def parseArgs(argv):
	parser = argparse.ArgumentParser(prog="diploid_calling.py", description=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter)
	parser.add_argument("invcf", type=str, help="A vcf file generated from cuteSV.")
	parser.add_argument("outvcf", type=str, help="The output diploid based SV callsets.")
	args = parser.parse_args(argv)
	return args

def setupLogging(debug=False):
	logLevel = logging.DEBUG if debug else logging.INFO
	logFormat = "%(asctime)s [%(levelname)s] %(message)s"
	logging.basicConfig( stream=sys.stderr, level=logLevel, format=logFormat )
	logging.info("Running %s" % " ".join(sys.argv))

if __name__ == '__main__':
	main(sys.argv[1:])

