#!/home/asiimwe/anaconda3/bin/python
#Extracting allele specific information per position overlapped by a read
import sys
import csv
import os
import string
import subprocess
from itertools import chain
import pandas as pd
from os import mkdir
import shutil
import glob
import fileinput
import stat
from subprocess import check_call
from subprocess import PIPE, run, check_output
from pysam import VariantFile
import json
from copy import deepcopy

#pseudo: for sample in samples read vcf file; create dictionary for each vcf/sample based on chr, pos and sample - values == REF, ALT, GT; for each key in subset search dictionary

#Samples run in subset 1: "HG00512",  "HG00513",  "HG00731",  "HG00732",  "HG00733",  "NA19238",  "NA19239",  "NA19240"
 
#samples = ["HG00512",  "HG00513",  "HG00731",  "HG00732",  "HG00733",  "NA19238",  "NA19239",  "NA19240"]

samples = ["HG00512"]

#Loading file with reads(with consistent haplotypes) that overlap variants
file_subset1 = "/scratch/asiimwe/ref_bias_STAR_WASP/samples_subset1.txt"
f = open(file_subset1, "r")

#sample_dict = "/scratch/asiimwe/ref_bias_STAR_WASP/samples_dict_subset1_chr1.txt"
#x = open(sample_dict, "a")


#mapped_subset1 = "/scratch/asiimwe/ref_bias_STAR_WASP/mapped_subset1.txt"
mapped_subset1 = "/scratch/asiimwe/ref_bias_STAR_WASP/mapped_subset1.txt"
file_out = open(mapped_subset1, "a")

for s in samples:
	vcf_path = "/scratch/asiimwe/STAR-WASP_FASTQs_VCFs/VCF/" + s + ".vcf"
	#print(vcf_path)
	
	#creating search dictionaries from each vcf file
	vcf_file = VariantFile(vcf_path)
	for line in vcf_file.fetch():
		pos='%s' %(line.pos)
		#print(line2)
		chr='%s' %(line.contig)
		#id=line.id - we have none
		ref='%s' %(line.ref)
		alt='%s' %(line.alts)
		#nsamples = len(line2.samples)
		#print(nsamples)
		gt = line.samples[0]['GT'] #GT is a 2 item tuple = (1,0)
		gt1 = '%s' %(gt[0])
		gt2 = '%s' %(gt[1])
		GT = gt1 + "|" + gt2
		#print(chr, " " , pos, " ", ref, " ", alt, " ", GT) #interested in these for the dict

		if len(ref) == 1 and len(alt) ==1:
			if GT == "1|0" or GT == "0|1":
				#print(ref, " ", alt)
				#print(chr, " " , pos, " ", ref, " ", alt, " ", GT)
				#vars = [chr, pos, ref, alt, GT]
				#for i in vars:
				#	locals()[i] = dict()
				#	print(dict.keys())
				d = {(chr, pos, s): [ref, alt, GT]}
				#d = dict()
				#for keys in d.keys():
				#	print(keys)
				keys = d.keys()
				values = d.values()
				#print(s, keys, values)
				#dline =  keys + " " +values
				#print(d)
				#x.write(str(d))
	
				#for key, value in d.items(): 
				#	x.write('%s:%s\n' %(key, value))


for i, inputline in enumerate(f): # Reading lines in sam file extract. Note that our search key will be based on the vG position(+1) to map positions in the vcf file - the vG coordinates are 0-based, while VCF are 1-based so we add 1 to vG values to find matching positions in the VCF file 
	#print(inputline) #"ERR1050076.12463822" "chr6" 32443897 "vA:B:c,1,1,1" "vG:B:i,32444793,32444802,32444814" "32444793" 32444794 "vW:i:1" "HG00512" "/scratch/asiimwe/WASP_FASTQs_VCFs/VCF/HG00512.vcf"
	#splitting line by space to extract individual elements
	read_id = '%s' %(inputline.split(" ")[0])
	chr = '%s' %(inputline.split(" ")[1])
	sam_pos = '%s' %(inputline.split(" ")[2])
	vA = '%s' %(inputline.split(" ")[3])
	vG = '%s' %(inputline.split(" ")[4])
	var_vG_pos = '%s' %(inputline.split(" ")[5]) #Overall we want to take the coordinates at which the reads overlap variants and search for these in the vcf file to extract specific information on overallaped variant 
	var_vG_pos_plusone ='%s' %(inputline.split(" ")[6])
	vW_Tag = '%s' %(inputline.split(" ")[7])
	sample = '%s' %(inputline.split(" ")[8])
	path = '%s' %(inputline.split(" ")[9])
	sam_search_key = (chr, var_vG_pos_plusone, sample)
	#print(sam_search_key)
	#print(read_id, " ", chr, " ", sam_pos, " ", vA, " ",vG, " ", var_vG_pos, " ", var_vG_pos_plusone," ", vW_Tag, " ", sample, " ", path, " ", sam_search_key) #verifying that we are extracting all items that may be needed downstream
	#"ERR1050076.13240386"   "chr5"   176346147   "vA:B:c,1"   "vG:B:i,176346283"   "176346283"   176346284   "vW:i:1"   "HG00512"   "/scratch/asiimwe/WASP_FASTQs_VCFs/VCF/HG00512.vcf"
	#creating a sam file-based search key that will be used to search through the vcf dictionary
	#[value for key, value in dictionary.items() if sam_search_key in key]
	#def search(c):
	#       for k, v in dictionary.items():		
	#               if type(k) in [list, tuple, dict] and c in k:
	#                       return c, k, v
	#               elif c == k:
	#                       return c, k, v
	#search(sam_search_key)
	#print(type(sam_search_key))
	#print(type(dictionary.keys()))
	#print(sam_search_key)	
	#return(read_id, chr, sam_pos, vA, vG, var_vG_pos,var_vG_pos_plusone,var_vG_pos_plusone, sample, path, sam_search_key)
	#for sam_search_key in dictionary:
	#       for k, v in dictionary.items():
	#               if sam_search_key == k:
	#                       line = read_id + "\t" + chr + "\t" + sam_pos + "\t" + sam_search_key + '%s' %(k)  + "\t" +'%s' %(v) + "\t" +  vA + "\t" + vG + "\t" + var_vG_pos + "\t" + var_vG_pos_plusone + "\t" + vW_Tag + "\t" + sample + "\t" + path
	#                       line2 = sam_search_key + '%s' %(k)  + "\t" +'%s' %(v)
	#                       print(line2)
	#                       #file_out.write(line2)    

	for sam_search_key in d:
		line = "{}".format(str(sam_search_key)) + "\t" +  '%s' %(d[sam_search_key]) + "\n"
		print(line)
		#file_out.write(line)