#!/usr/bin/env python

from __future__ import print_function
import os
import sys
import gzip
import argparse
import csv
import json
import re
from collections import defaultdict
from collections import OrderedDict

parser = argparse.ArgumentParser()
parser.add_argument("vcf_dir", help="Directory for input vcf files generated by lumpy")
args = parser.parse_args()

count = 0

valid_chr = range(1, 23, 1)
valid_chr = [str(item) for item in valid_chr ]
valid_chr.append('X')
valid_chr.append('Y')

vcf_header = ["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER", "INFO", "FORMAT", "DATA"]
annotation_field = ["Consequence", "Codons", "Amino_acids", "Gene", "SYMBOL", "Feature", "EXON", "PolyPhen", "SIFT", "Protein_position", "BIOTYPE"]

features = {"CHROM": str, "POS": int, "ID": str, "REF": str, "ALT":str, "QUAL": float, "FILTER": str, "SVTYPE": str, "SVLEN": int, "END": int, "STRANDS": str, "IMPRECISE": str, "CIPOS": str, "CIEND": str, "CIPOS95": str, "CIEND95": str, "MATEID": str, "EVENT": str, "SECONDARY": str, "SU": int, "PE": int, "SR": int, "EV": str, "PRPOS": str, "PREND": str, "GT": str, "BD": int, "GQ": int, "SQ": float, "GL": str, "DP": int, "RO": int, "AO": int, "QR": int, "QA": int, "RS": int, "AS": int, "ASC": int, "RP": int, "AP": int, "AB": float, "CN": float}

# get all the vcf files in the input directory
vcfs = [ vcf for vcf in os.listdir(args.vcf_dir) if vcf.endswith('vcf.gz')]

# loop through each vcf file and extract data
for vcf in vcfs:
	sample = vcf.split("/")[-1].split('_')[0]
	vcf = os.path.join(args.vcf_dir, vcf)	
	
	with gzip.open(os.path.abspath(vcf), 'rt') as f:
		for line in f.readlines():              
			tmp_data = OrderedDict()
			
			# initialize an empty dictionary
			for key in features.keys():
				tmp_data[key] = None
			
			if not line.startswith('#'):
				tmp = line.strip().split("\t")
	
				for i in range(len(vcf_header)):
					tmp_data[vcf_header[i]] = tmp[i]
	
				if not tmp[0] in valid_chr:
					continue
	
				try:	
					info_fields = tmp[7].split(";")
				except ValueError as e:
					print(e, file=sys.stderr)
					continue
			
				for item in info_fields:
					if item == 'IMPRECISE':
						item = 'IMPRECISE=true'
					if item == 'SECONDARY':
						item = 'SECONDARY=true'
					try:
						key, value = item.split("=")
						tmp_data[key] = value
					except ValueError:
						print(e, file=sys.stderr)
						continue
				
				sv_name = tmp[0] + '_' + tmp[1] + tmp_data['SVTYPE']
				tmp_data['sv_name'] = sv_name
				
				# create a dict to store format and format data
				format_ = tmp_data['FORMAT'].split(':')
				value_ = tmp_data['DATA'].split(':')
	
				for i in range(len(format_)):
					tmp_data[format_[i]] = value_[i]
				
				del tmp_data['INFO']
				del tmp_data['FORMAT']
				del tmp_data['DATA']
				#del tmp_data['ALT']
				
				# convert to desired data type
				for key in features:
					if tmp_data[key] == '.':
						tmp_data[key] = None	
					if tmp_data[key] is not  None:	
						try:		
							tmp_data[key] = features[key](tmp_data[key])

						except (ValueError, TypeError) as e:
							#print >> sys.stderr
							print(e, file=sys.stderr)
							continue
						
				tmp2 = []
				
				if 'CSQ' in tmp_data.keys():	
					annotations = tmp_data['CSQ'].split(',')
					del tmp_data['CSQ']
					
					for anno in annotations:
						tmp3 =  anno.split('|')
						tmp4 = defaultdict()
						for i in range(len(tmp3)):
							tmp4[annotation_field[i]] = tmp3[i] 		
								
						tmp2.append(tmp4)
				
				tmp_data['annotation'] = tmp2			


				count += 1
				tmp_data['sample'] = sample
				
				print ('{"index":{"_id": "', count, '"}}', sep='')           
				print(json.dumps(tmp_data))

