#!/usr/bin/env python2

import sys,re

from bisect import bisect_right

ref_dict={}
alt_dict={}


# read in snps.tsv, put in ref and alt dictionaries
snpfile=open(sys.argv[1])
for l in snpfile:
	split=l.split()
	if split[0] in ref_dict:
		ref_dict[split[0]][int(split[1])]=split[2]
		alt_dict[split[0]][int(split[1])]=split[3]
	else:
		ref_dict[split[0]]={}
		alt_dict[split[0]]={}
		ref_dict[split[0]][int(split[1])]=split[2]
		alt_dict[split[0]][int(split[1])]=split[3]

snpfile.close()

chrs=list(ref_dict.keys())
chrs.sort()
key_dict={}

# I believe this creates key_dict, which stores the positions of all snps on each chrom, in order
for c in chrs:
	keys=list(ref_dict[c].keys())
	keys.sort()
	key_dict[c]=keys
		
# reads in the bam file
for l in sys.stdin:
    # prints out the header
	if l[0] == '@':
		print l.rstrip()
		continue
	split=l.rstrip().split('\t')
	read=split[9]
	readlen=len(read)
	ch=split[2]
	bp=int(split[3])

	md=l.rstrip().split('MD:Z:')[1].split('\t')[0]
    # if deletion in the string relative to the reference, don't consider this read
	if md.find('^')>0:
		continue
    # split by non-digits (i.e., AGCT) to get the parts that match 
	matches=re.split('\D+',md)
    # split by digits to get the parts that don't match
	mismatches=re.split('\d+',md)
	s=md
	ref=''
	i=0
	while len(s)!=0:
		i=len(ref)
        # if md starts with one or more AGCT, take those letters and append them to the reference string
		if re.findall('\A\D+',s):
			x=re.findall('\A\D+',s)[0]
			ref+=x
			s=re.sub('\A\D+','',s)
		
        # if md starts with one or more numbers, take those numbers and append those positions in the read to the reference string
		elif re.findall('\A\d+',s):
			x=int(re.findall('\A\d+',s)[0])
			s=re.sub('\A\d+','',s)
			ref+=read[i:(i+x)]

    # if this all worked correctly (does this ever not work? maybe in the case of insertions) now you have the reference sequence. But what we really want is the alternative sequence. 
	if len(ref)==len(read):

		start=bp
		end=bp+readlen
		keys=key_dict[ch]
		read=list(read)
		index=bisect_right(keys,start)-1
		split[0]="mod_"+split[0]
        
        # this takes in the sequence, changes ref to alt and alt to ref. But a bit odd because to get the read, it seems like you're only considering the reference. 
		for i in xrange(index,len(keys)):
			k=keys[i]
			if k<start:
				continue
			elif k>=start and k<end:
				snp_index=k-start		
				if read[snp_index]==ref_dict[ch][k]:
					read[snp_index]=alt_dict[ch][k]
				elif read[snp_index]==alt_dict[ch][k]:
					read[snp_index]=ref_dict[ch][k]
							
			elif k>end:
				break
		split[9]=''.join(read)
		print '\t'.join(split)
    #else:
     #   print('ERROR, ref and alt reads are not the same length!')
