Workflow Steps and Code Snippets

from sys import argv
from os import path, SEEK_CUR
from scipy.stats import norm
from collections import defaultdict
import argparse
import csv
import pysam

import pdb

genome_length = int(3e9)

def main(argv):
	#get arguments
	parser = argparse.ArgumentParser(description='simulate viral integrations')
	parser.add_argument('--sim-info', help='information from simulation', required = True, type=str)
	parser.add_argument('--sim-bam', help='bam file from read simulation (must be sorted, indexed)', required = True, type=str)
	parser.add_argument('--soft-threshold', help='threshold for number of bases that must be mapped when finding integrations in soft-clipped reads', type=int, default=20)
	parser.add_argument('--discordant-threshold', help='threshold for number of bases that may be unmapped when finding integrations in discordant read pairs', type=int, default=20)
	parser.add_argument('--mean-frag-len', help='mean framgement length used when simulating reads', type=int, default=500)
	parser.add_argument('--sd-frag-len', help='standard devation of framgement length used when simulating reads', type=int, default=30)
	parser.add_argument('--window-frac', help='fraction of the distribution of fragment sizes to use when searching for discordant read pairs', type=float, default=0.99)
	parser.add_argument('--output', help='output file', required=False, default='annotated.tsv')
	args = parser.parse_args()

	# read in bam file
	samfile = pysam.AlignmentFile(args.sim_bam)

	# iterate over integrations in info file and pull out reads crossing each one
	# the sam/bam file has the same coordinates as the info file
	# so just use the coordinates of the left and right junctions from this files
	with open(args.sim_info, newline='') as info_file, open(args.output, 'w', newline='') as output:

		# use csv to read and write files
		reader = csv.DictReader(info_file, delimiter = '\t')

		writer_fieldnames = list(reader.fieldnames) + ['left_chimeric', 'right_chimeric', 'left_discord', 'right_discord', 'multiple_discord', 'fake_discord', 'discordant_and_chimeric']

		writer = csv.DictWriter(output, delimiter = '\t', fieldnames = writer_fieldnames)
		writer.writeheader()


		# set window size for looking for discordant pairs and 
		# looking for multiple ints with the same discordant pair
		window_width = window_size(args.mean_frag_len, args.sd_frag_len, args.window_frac)

		# create a buffer of integrations that all fall within this window width 
		buffer = [next(reader)]
		while True:

			buffer, next_buffer = get_ints_in_window(buffer, window_width, reader, info_file)

			# find reads crossing each integration in the buffer
			buffer = find_reads_crossing_ints(buffer, samfile, args, window_width)

			# check for read pairs that cross multiple junctions
			buffer = find_multiple_discordant(buffer)

			# check for read pairs that are discordant at one integration and chimeric at another integration
			# assume that integrations are independent (even though they aren't in simulation), so these are not detectable
			buffer = find_multiple_chimeric_discordant(buffer)

			# write rows in buffer
			for row in buffer:
				writer.writerow(row)

			buffer = next_buffer

			# check if we've reached the end of the file
			if len(buffer) == 0:
				break

	samfile.close()

def find_multiple_chimeric_discordant(buffer):
	"""
	Assume that in real data, integrations are independent, so two integrations do not
	occur close enough to each other in the same cell that we can detect them.  This may
	or may not be the case, but the number of these kinds of pairs of integrations is assumed
	to be small enough that they're not worth worrying about. 

	If a read pair has one read that is chimeric, but it's also discordant about the same
	or a different integration, only the chimeric read will be detected 
	(and it won't be flagged as a discordant read pair). Assuming that
	these events are a tiny minority in real data, here don't flag the discordant pairs
	as something that should be detected (but instead add them to the 'discordant_and_chimeric'
	category)
	"""	
	to_delete = {}
	for i, row in enumerate(buffer):

		# for each row, find reads that are both chimeric
		# and discordant for any other integration
		left_discord = row['left_discord'].split(';')
		right_discord =  row['right_discord'].split(';')
		to_delete[i] = {'left' : [], 'right' : []}


		# check for these reads in other rows
		for row_2 in buffer:
			# skip if this is the same as the one we're looking at
			if row_2 == row:
				continue

			# get read names for chimeric reads for this row
			chimeric = row_2['left_chimeric'].split(";")
			chimeric += row_2['right_chimeric'].split(";")
			chimeric = [read[:-2] for read in chimeric]

			# check if chimeric reads in row_2 also occur in the list
			# of discordant reads in row
			for read in chimeric:
				if read in left_discord:
					to_delete[i]['left'].append(read)
				if read in right_discord:
					to_delete[i]['right'].append(read)

	# remove reads in to_delete
	for row_num in to_delete.keys():

		# remove reads from left_discord
		left_discord = buffer[row_num]['left_discord'].split(';')
		left_discord = [read for read in left_discord if read not in to_delete[row_num]['left']]
		buffer[row_num]['left_discord'] = ";".join(left_discord)

		# remove reads from right_discord
		right_discord = buffer[row_num]['right_discord'].split(';')
		right_discord = [read for read in right_discord if read not in to_delete[row_num]['right']]
		buffer[row_num]['right_discord'] = ";".join(right_discord)		

		# removed reads go to 'discordant_and_chimeric'
		buffer[row_num]['discordant_and_chimeric'] = [f"{read}_left" for read in to_delete[row_num]['left']]
		buffer[row_num]['discordant_and_chimeric'] += [f"{read}_right" for read in to_delete[row_num]['right']]
		buffer[row_num]['discordant_and_chimeric'] = ";".join(buffer[row_num]['discordant_and_chimeric'])

	return buffer

def find_multiple_discordant(buffer):
	"""
	a read might be discordant about two or more integration sites
	for example, if a read pair crosses the right border of one integration and the left
	border of a close by integration, both read pairs will be mapped to the virus
	and therefore this pair won't be detected as an integration

	deal with discordant pairs that cross multiple integrations in one of two ways:
		1. if a pair is discordant about the right side of one integration and the left side
		of a close by integration, then both reads are mapped to the virus and this read should
		be removed from both.  Keep track of these in a different category ('fake_discord')
		2. if a pair crosses right side of one integration, and both sides of a nearby integration,
		then one read maps to virus (in the first integration), and the other maps to host (after the second)
		and therefore this pair should be indicated to be an integration.  However, if we want to score based
		on integration position, we need to decide which integration this read 'belongs' to.  Keep
		these reads in a seperate category ('multiple_discord') for all the integrations that they cross,
		to reflect the fact that the 'belong' to multiple integrations
	"""
	# keep track of which reads we've already seen
	seen = dict()

	# for reads that we find at multiple junctions, 
	# keep track of which junctions so we can work out what to do with that read
	multiples = dict()

	for i, row in enumerate(buffer):
		# check reads found around left side
		for read in row['left_discord'].split(';'):
			if read == '':
				continue
			if read in seen.keys():
				# if this is the first double, add the entry from seen
				if read not in multiples:
					multiples[read] = []
					multiples[read].append(seen[read])

				# add entry to multiples for this time we found the read
				multiples[read].append({'int_id' : row['id'], 'side' : 'left' , 'buffer_row' : i})

			else:
				seen[read] = {'int_id' : row['id'], 'side' : 'left', 'buffer_row' : i}
		# check right side
		for read in row['right_discord'].split(';'):
			if read == '':
				continue
			if read in seen.keys():
				# if this is the first double, add the entry from seen
				if read not in multiples:
					multiples[read] = []
					multiples[read].append(seen[read])

				# add entry to multiples for this time we found the read
				multiples[read].append({'int_id' : row['id'], 'side' : 'right', 'buffer_row' : i})
			else:
				seen[read] = {'int_id' : row['id'], 'side' : 'right', 'buffer_row' : i}


		# add extra keys to dicts 
		row['multiple_discord'] = []
		row['fake_discord'] = []

	# deal with reads crossing multiple integrations
	for read in multiples:

		# if the side of the first junction matches the side of the last junction
		# this is a 'multiple_discord'
		sides = [find['side'] for find in multiples[read]]
		if sides[0] == sides[-1]:
			new_type = 'multiple_discord'
		else:
			new_type = 'fake_discord'

		for find in multiples[read]:

			# find row and side
			buffer_row = find['buffer_row']
			old_type = f"{find['side']}_discord"

			# remove from list of reads of old type
			reads = buffer[buffer_row][old_type].split(";")
			reads.remove(read)
			buffer[buffer_row][old_type] = ";".join(reads)

			# add to list of new type
			buffer[buffer_row][new_type].append(f"{read}_{find['side']}")

	# for each row, join lists of 'multiple_disord' and 'fake_discord'
	for row in buffer:
		row['multiple_discord'] = ";".join(row['multiple_discord'])
		row['fake_discord'] = ";".join(row['fake_discord'])

	return buffer

def times_already_found(read, multiples):
	"""
	check how many times a read has already been found when checking for discordant about multiple integration sides
	"""	
	if read in multiples:
		return len(multiples[read])
	else:
		return 0

def find_reads_crossing_ints(buffer, samfile, args, window_width):
	"""
	find reads crossing the integration site, and add them to the row
	reads crossing the integration site can be chimeric about the left or right junction
	or discordant about the left or right junction
	if they're discordant about both junctions, they actually go from host sequence to host sequence
	and	therefore aren't actually discordant
	"""
	for row in buffer:

		# get information about integration site location
		chr = row['chr']
		left_start = int(row['leftStart']) 
		left_stop = int(row['leftStop'])
		right_start = int(row['rightStart']) 
		right_stop = int(row['rightStop'])

		assert left_start >= 0 and right_start >= 0

		# find chimeric reads
		left_chimeric = get_chimeric(chr, left_start, left_stop, samfile, args.soft_threshold, buffer)
		right_chimeric = get_chimeric(chr, right_start, right_stop, samfile, args.soft_threshold, buffer)

		row['left_chimeric'] = ";".join(left_chimeric)
		row['right_chimeric'] = ";".join(right_chimeric)

		# find discordant read pairs	
		left_discord = get_discordant(chr, left_start, left_stop, samfile, args.discordant_threshold, window_width, buffer)
		right_discord = get_discordant(chr, right_start, right_stop, samfile, args.discordant_threshold, window_width, buffer)

		# if a read is both chimeric and discordant, chimeric takes priority 
		# (but it shouldn't be both if the clipping threshold is the same for both read types)
		left_chimeric, left_discord = remove_chimeric_from_discord(left_chimeric, left_discord)
		right_chimeric, right_discord = remove_chimeric_from_discord(right_chimeric, right_discord)

		left_chimeric, right_discord = remove_chimeric_from_discord(left_chimeric, right_discord)
		right_chimeric, left_discord = remove_chimeric_from_discord(right_chimeric, left_discord)


		row['left_discord'] = ";".join(left_discord)
		row['right_discord'] = ";".join(right_discord)

	return buffer	

def get_ints_in_window(buffer, window_width, reader, reader_handle):
	"""
	get a list of integrations (dicts corresponding to one row from the int-info file)
	that are within window_width of each other
	"""	

	assert len(buffer) == 1

	# get position and chromosome from buffer
	start = int(buffer[0]['rightStop'])
	chr = buffer[0]['chr']

	# get more integraions until we're not in the window anymore or we're at the end of the file
	# to avoid having to go back a line, save the first line not added to 'next_buffer'
	while True:	

		# get next row
		try:
			row = next(reader)
		except StopIteration:
			next_buffer = []
			break

		# compare previous integration with this integration
		prev_start = start
		prev_chr = chr
		start =  int(row['leftStart'])
		chr = row['chr']

		# check if next row is a window width away
		if (start - prev_start < window_width) and prev_chr == chr:
			# add to buffer
			buffer.append(row)
			start = int(row['rightStop'])
		else:
			# don't add the row but this will be the start of the next buffer
			next_buffer = [row]
			break

	return buffer, next_buffer	

def get_discordant(chr, start, stop, samfile, threshold, window_width, buffer):
	"""
	Get any discordant read pairs which cross an integration site
	In other words, get pairs where one mate is mapped on the host side, and the other on the virus side
	A read is considered mapped if it has at most threshold (default 20) unmapped bases

	This includes any pairs where the integration site falls within threshold (default 20) bases 
	of the end of the read (for a read mapped on the left of the integration), 
	or within threshold bases of the start of the read for a read on the mapped on the right
	of the integration

	Avoid an exhaustive search by extracting only the reads in a window around the integration site
	Set this window based on the mean length and standard deviation of fragment size used in simulation
	and the fraction of the fragment length distribution we want to cover.  Find discordant
	pairs by finding pairs for which an integration (start, stop) falls between the 20th
	base from the end of the left read and the 20th base of the right read

	Current criteria in discordant.pl is that a pair must have one read mapped and the other
	unmapped to be considered a discordant read pair.  To be considered 'mapped', a read must
	have 20 or fewer unmapped bases.  To be considered 'unmapped', a read must have
	less than 20 unmapped bases

	Note that a read pair might have one integration fall between read1 and read2, but read1 or read2
	might also cross a second integration.  This pair is therefore both discordant about one integration, and
	also one member of the pair is chimeric  A consequence of this is that one read maps only to vector or host, 
	but the other maps to both.  This discordant pair cannot currently be detected, since
	the pipeline currently detects discordant read-pairs only if one read is mapped to vector and
	not host, and vice versa for the other read.  However, this pair really is evidence
	for integration at the first site (as a discordant read-pair), so include it in the
	output as such.
	"""

	reads = []

	# extract read pairs in the desired window
	# pysam numbering is 0-based, with the only exception being the region string in the fetch() and pileup() methods. 
	window_start = start - window_width
	if window_start < 1:

		window_start = 1

	window_stop = stop + window_width
	if window_stop > samfile.get_reference_length(chr):
		window_stop = samfile.get_reference_length(chr) 
	if window_stop == window_start:
		window_stop += 1

	for read1, read2 in read_pair_generator(samfile, f"{chr}:{window_start}-{window_stop}"):

		# check mate is mapped
		if read1.is_unmapped or read2.is_unmapped:
			continue
		# check reference for this read is the same as mate
		if read1.reference_name != read2.reference_name:
			continue
		# if this read is forward, mate must be reverse and vice versa
		if (read1.is_reverse == read2.is_reverse):
			continue

		# if the integration site falls between left_boundary and right_boundary
		# (which are coordinates within the reference)
		# this pair crosses the integration site
		if read1.is_reverse is False:
			left_boundary = get_boundary(read1, threshold, side = "right")
			right_boundary =  get_boundary(read2, threshold, side = "left")
		else:
			left_boundary = get_boundary(read2, threshold, side = "right")
			right_boundary =  get_boundary(read1, threshold, side = "left")

		# if left_boundary is greater than right_boundary, reads overlap
		if left_boundary >= right_boundary:
			continue

		assert left_boundary is not None
		assert right_boundary is not None
		assert left_boundary < right_boundary

		if within(start, stop, left_boundary, right_boundary):

			reads.append(read1.qname)
# 			
# 			# TODO - need to decide if integrations should be excluded on the basis
# 			# below (if it's not the case that one read is mapped to host and the other to virus)
# 			# these events can't be currently detected in the pipeline, but are real evidence
# 			# of integration, so for now include them in the output
# 
# 			r1_mapped = get_mapped_ref(read1, buffer, threshold)
# 			r2_mapped = get_mapped_ref(read2, buffer, threshold)
# 			
# 			assert r1_mapped['host'] or r1_mapped['virus']
# 			assert r2_mapped['host'] or r2_mapped['virus']
# 			
# 			if r1_mapped['host'] != r2_mapped['host'] and r1_mapped['virus'] != r2_mapped['virus']:
# 				reads.append(read1.qname)

	return reads

def get_mapped_ref(read, buffer, threshold):
	"""
	figure out if each read in this read pair will be mapped to host or vector/virus
	returns a dict with the keys 'host' and 'virus',
	and values True or False depending on if the read is mapped to either or not
	"""

	assert read.is_unmapped is False

	read_mapped = {'host':False, 'virus':False, 'int' : []}

	# get first and last position to which read is mapped in reference
	first = read.get_reference_positions()[0]
	last =  read.get_reference_positions()[-1]

	prev_start = 0
	for row in buffer:
		# figure out if we need to include the ambiguous bases or not
		if row['juncTypes'].split(',')[0] == 'gap':
			left_host_junc = int(row['leftStart'])
			left_virus_junc = int(row['leftStop'])
		else:
			left_host_junc = int(row['leftStop'])
			left_virus_junc = int(row['leftStart'])
		if row['juncTypes'].split(',')[1] == 'gap':
			right_virus_junc = int(row['rightStart'])
			right_host_junc = int(row['rightStop'])
		else:
			right_virus_junc = int(row['rightStop'])
			right_host_junc = int(row['rightStart'])	

		# if read is between the start of the chromosome and the start of the integration
		if intersects(first, last, prev_start, left_host_junc):
			# check we have at least threshold bases mapped
			if check_threshold(read, prev_start, left_host_junc, threshold):
				read_mapped['host']  = True
		# if read is between the left and right junctions
		if intersects(first, last, left_virus_junc, right_virus_junc):
			if check_threshold(read, left_virus_junc, right_virus_junc, threshold):
				read_mapped['virus'] = True
				read_mapped['int'].append(row['id'])
		prev_start = right_host_junc

		# if the prev_start is to the right of the position of the last mapped base in the read
		# then we don't need to continue
		if prev_start > last:
			break


	# check if read is between the end of the last integration and the end of the chromosome
	# don't know what the end of the chromosome is, so just use a large number (genome_length)
	if intersects(first, last, prev_start, genome_length):
		if check_threshold(read, prev_start, genome_length, threshold):
			read_mapped['host'] = True

	return read_mapped	

def get_chimeric(chr, start, stop, samfile, threshold, buffer):
	"""
	find reads that cross an interval defined as chr:start-stop in samfile
	the interval must be at least threshold bases from the start and end of the read
	"""
	reads = []

	# get reads that cross interval
	# pysam numbering is 0-based, with the only exception being the region string in the fetch() and pileup() methods. 
	# The same is true for any coordinates passed to the samtools command utilities directly, such as pysam.fetch().
	for read in samfile.fetch(chr, start, stop + 1):

		# check that interval is at least threshold bases from either end of the read
		mapped = get_mapped_ref(read, buffer, threshold)
		assert (mapped['host'] or mapped['virus'])

		if mapped['host'] is True and mapped['virus'] is True:

			reads.append(read.query_name + read_num(read))

	return reads

def read_num(read):
	"""
	return '/1' if read is R1, '/2' if read is R2, or empty string otherwise
	"""
	if read.is_read1 is True:
		return "/1"
	elif read.is_read2 is True:
		return "/2"
	else:
		return ""

def check_threshold(read, start, stop, threshold):
	""""
	check that there are least threshold bases that map to an interval (defined by start and stop)
	"""
	rstart = read.get_reference_positions()[0]
	# need to account for 0-based numbering (stop already has this accounted for)
	rstop = read.get_reference_positions()[-1] + 1
	assert intersects(rstart, rstop, start, stop)

	if (rstop - start) < threshold:
		return False
	if (stop - rstart) < threshold:
		return False

	return True

def window_size(mean_frag_len, sd_frag_len, window_frac):
	"""
	to avoid exhaustive search for discordant reads, 
	work out window size for extracting reads when finding discordant read pairs
	based on mean and standard deviation of fragment length distribution,
	and fraction of this distribution we want to cover

	For example, if the mean fragment length is 500 bp, the standard deviation is 30 bp, and we want 
	to cover 0.99 of this distribution, the window size would be 570 bp, which accounts for 99% of 
	the fragment sizes in this distribution (one-tailed)

	"""

	# get one-tailed value which contains window_frac of fragments
	upper = norm.ppf(window_frac, loc = mean_frag_len, scale = sd_frag_len)

	return int(round(upper))

def get_boundary(read, threshold, side = "left"):
	"""
	get first position in reference that is at least threshold bases from the 
	start or end of the read and isn't None

	if side is 'left', return the 0-based position after threshold mapped bases from the 
	start of the read
	if side is 'right', return the 0-based position before threshold mapped bases from the
	end of the read
	"""
	assert isinstance(threshold, int)
	assert threshold >= 0
	assert threshold <= read.qlen
	assert side in ['left', 'right']

	aligned_pairs = read.get_aligned_pairs()

	if side == "left":
		# if we want to look on the right hand side, we need to look backwards 
		# through the aligned pairs
		aligned_pairs = list(reversed(aligned_pairs))
		# python numbering is zero-based and half-open
		threshold -= 1 


	# iterate over bases in aligned_pairs, starting from the threshold value
	for pair in aligned_pairs[threshold:]:
		# check the base is aligned
		if (pair[1] is not None):
			return pair[1]

def within(start1, stop1, start2, stop2):
	"""
	compare two intervals, each with a start and stop value
	return true if the first interval is completely within in the second

	use half-open intervals, so [8, 8) is not within [5, 8) 
	within(8, 8, 5, 8) => False
	within(6, 6, 5, 8) => True
	within(5, 8, 5, 8) => True
	within(4, 6, 5, 8) => False
	within(5, 9, 5, 8) => False
	"""	
	assert start1 <= stop1
	assert start2 <= stop2
	assert isinstance(start1, int)
	assert isinstance(start2, int)
	assert isinstance(stop1, int)
	assert isinstance(stop2, int)

	if start1 >= start2 and start1 < stop2:
		if stop1 - 1 >= start2 and stop1 - 1 < stop2:
			return True
	return False

def intersects(start1, stop1, start2, stop2):
	"""
	compare two intervals, each with a start and stop value
	return true if the first interval is intersects the second

	use half-open intervals, so [2, 3) and [3, 4) don't intersect
	"""	
	assert start1 <= stop1
	assert start2 <= stop2
	assert isinstance(start1, int)
	assert isinstance(start2, int)
	assert isinstance(stop1, int)
	assert isinstance(stop2, int)

	# if first interval is completely to the left of the second interval
	if stop1 <= start2:
		return False
	# if first interval is completely to the right of the second interval
	if start1 >= stop2:
		return False

	# otherwise, they intersect
	return True	

def read_pair_generator(bam, region_string=None):
    """
    https://www.biostars.org/p/306041/
    Generate read pairs in a BAM file or within a region string.
    Reads are added to read_dict until a pair is found.
    """
    read_dict = defaultdict(lambda: [None, None])
    for read in bam.fetch(region=region_string):
        if read.is_secondary or read.is_supplementary:
            continue
        qname = read.query_name
        if qname not in read_dict:
            if read.is_read1:
                read_dict[qname][0] = read
            else:
                read_dict[qname][1] = read
        else:
            if read.is_read1:
                yield read, read_dict[qname][1]
            else:
                yield read_dict[qname][0], read
            del read_dict[qname]

def remove_chimeric_from_discord(chimeric, discord):
	"""
	check for read ids that are in both chimeric and discord - remove them from discord if they're in both
	"""
	# chimeric reads have /1 or /2 added
	to_delete = []
	chimeric_tmp = [read[:-2] for read in chimeric]
	for read in discord:
		if read in chimeric_tmp:
			print(f"  removed a read that was in both chimeric and discord: {read}")
			to_delete.append(read)

	# remove reads flagged for deletion
	discord = [read for read in discord if read not in to_delete]

	return chimeric, discord

if __name__ == "__main__":
	main(argv[1:])

import csv
import os, glob, shutil
import re
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from qmplot import manhattanplot
import natsort
from natsort import natsorted
import seaborn as sns
import pysam

if __name__ == "__main__":

    ## Logging
    with open(snakemake.log[0], "w") as f:
      sys.stderr = sys.stdout = f

      ## Read the alignment data 
      align_kmers_sam = pd.read_table(snakemake.input[0], 
                                      sep='\t', comment='@', header=None,
                                      usecols=[0,2,3], names=['kmer_id', 'chr', 'bp'])

      ## Preparing the data for plotting
      align_kmers_sam['kmer'] = align_kmers_sam['kmer_id'].str.split('_').str[0]
      align_kmers_sam['p_value'] = align_kmers_sam['kmer_id'].str.split('_').str[1]
      align_kmers_sam['p_value'] = align_kmers_sam['p_value'].astype(float)
      align_kmers_sam['bp'] = align_kmers_sam['bp'].astype(int)

      # Sort the data by chromosome and chromosome position
      align_kmers_sam_sorted = align_kmers_sam.sort_values(by=["chr", "bp"], key=natsort.natsort_keygen())

      # Get colors for manhattan plot
      colors = sns.color_palette("colorblind").as_hex()
      colors_2 = sns.color_palette("husl").as_hex()

      # Make a column of minus log10 p-values
      align_kmers_sam_sorted['minuslog10pvalue'] = -np.log10(align_kmers_sam_sorted.p_value)

      ## Get min & max minus log10 p-values for y axis limits
      y_max = align_kmers_sam_sorted['minuslog10pvalue'].max()
      y_min = align_kmers_sam_sorted['minuslog10pvalue'].min()
      print("y_max: " + str(y_max))
      print("y_min: " + str(y_min))
      ## Check if only one chromosome is provided for the manhattan plot
      num_of_chrs = len(pd.unique(align_kmers_sam_sorted['chr']))

      # Define a function to extract chromosome names and lengths from a SAM file header
      def extract_chromosome_info(sam_file):
        """
        Extract chromosome names and lengths from a SAM file header and return as a Pandas DataFrame with columns "chr" and "bp".
        """
        chromosome_info = {} # dictionary to store chromosome names and lengths
        pattern = re.compile(r'^([Cc][Hh][Rr])?\d*[XYM]?$') # chromosome names pattern

        with pysam.AlignmentFile(sam_file, "r") as sam: # open SAM file
            for header_line in sam.header["SQ"]: # iterate over SQ header lines
                chromosome_name = header_line["SN"] # get chromosome name
                length = header_line["LN"] # get chromosome length
                if chromosome_name.startswith("chr"): # if chromosome name starts with "chr"
                    name = chromosome_name # use name as is
                else: # if chromosome name does not start with "chr"
                    match = pattern.match(chromosome_name) # 
                    if match: # if chromosome name matches pattern
                        # Convert name to "chrX" format
                        name = "chr" + match.group(1)
                    else: # if chromosome name does not match pattern
                        continue # skip chromosome
                chromosome_info[name] = length # add chromosome name and length to dictionary

        # Convert dictionary to DataFrame
        df = pd.DataFrame(chromosome_info.items(), columns=["chr", "bp"])

        return df # return the DataFrame

      # Plotting the manhattan plot
      print("Plotting...")

      # Set font sizes
      tick_fontsize = snakemake.params["tick_fontsize"]
      label_fontsize = snakemake.params["label_fontsize"]
      title_fontsize = snakemake.params["title_fontsize"]

      # Set the figure dpi
      dpi = snakemake.params["dpi"]

      ## If only one chromosome is provided, plot the k-mer's position on
      ## that chromosome on the x axis
      if num_of_chrs == 1:
        f, ax = plt.subplots(figsize=(18, 9), facecolor='w', edgecolor='k')
        manhattanplot(data=align_kmers_sam_sorted,
                      snp="kmer_id",
                      chrom="chr",
                      CHR= pd.unique(align_kmers_sam_sorted['chr']),
                      color=colors_2,
                      pos="bp",
                      pv="p_value",
                      suggestiveline=None,  # Turn off suggestiveline
                      genomewideline=None,  # Turn off genomewideline
                      xticklabel_kws={"rotation": "vertical"},
                      ax=ax,
                      s = snakemake.params["point_size"],
                      clip_on=False)
        ax.set_ylim([y_min-0.5, y_max+1]) # Set y axis limits

        # Set x axis tick interval
        xtick_interval = snakemake.params["xtick_interval"]

        # Calculate the minimum and maximum of your data, rounded to the nearest multiple of 5000
        min_val = align_kmers_sam_sorted['bp'].min() // xtick_interval * xtick_interval
        max_val = (align_kmers_sam_sorted['bp'].max() // xtick_interval + 1) * xtick_interval

        # Generate the tick locations
        xtick_locs = np.arange(min_val, max_val, xtick_interval)

        f.suptitle('k-mer Based GWAS Manhattan Plot for ' + snakemake.params["pheno"], fontsize=title_fontsize)
        plt.xlabel('Chromosome: ' + pd.unique(align_kmers_sam_sorted['chr'])[0], fontsize=label_fontsize)
        plt.ylabel(r"$-log_{10}{(P)}$", fontsize=label_fontsize) 
        plt.xticks(xtick_locs, fontsize = tick_fontsize)
        plt.yticks(fontsize = tick_fontsize)
        plt.tight_layout()

      ## If more than one chromosome is provided, use all chromosomes
      if num_of_chrs > 1:

        # Extract chromosome names and lengths from the SAM file header
        chrom_info_tab = extract_chromosome_info(snakemake.input[0])
        chrom_names = natsorted(chrom_info_tab['chr'].tolist())

        # Add extra chromosome names and lengths to the data frame
        align_kmers_sam_with_all_chrom = pd.concat([align_kmers_sam, chrom_info_tab], ignore_index=True)
        align_kmers_sam_with_all_chrom = align_kmers_sam_with_all_chrom[align_kmers_sam_with_all_chrom['chr'].isin(chrom_names)]

        # Sort the data by chromosome and chromosome position
        align_kmers_sam_with_all_chrom_sorted = align_kmers_sam_with_all_chrom.sort_values(by=["chr", "bp"], key=natsort.natsort_keygen()) 
        # Fill NaN values with 1
        align_kmers_sam_with_all_chrom_sorted['p_value'] = align_kmers_sam_with_all_chrom_sorted['p_value'].fillna(1) 

        # Plot the dots of chromosome length rows wiht the lowest opacity 
        extra_rows = align_kmers_sam_with_all_chrom_sorted[align_kmers_sam_with_all_chrom_sorted['p_value'] == 1]

        f, ax = plt.subplots(figsize=(18, 9), facecolor='w', edgecolor='k')
        manhattanplot(data=align_kmers_sam_with_all_chrom_sorted,
                      snp="kmer_id",
                      chrom="chr",
                      color=colors,
                      pos="bp",
                      pv="p_value",
                      suggestiveline=None,  # Turn off suggestiveline
                      genomewideline=None,  # Turn off genomewideline
                      xticklabel_kws={"rotation": "vertical"},
                      ax=ax,
                      s = snakemake.params["point_size"],
                      clip_on=True)
        ax.set_ylim([y_min-2, y_max+1]) # Set y axis limits
        plt.scatter(extra_rows['bp'], -np.log10(extra_rows['p_value']), alpha=0)

        # Calculate the cumulative distances from the start of each chromosome and store them in a list
        chrom_ends = chrom_info_tab['bp'].cumsum().tolist()

        # Plot the vertical lines for the end of each chromosome
        for end_position in chrom_ends:
          plt.axvline(x=end_position, color='grey', linestyle='--', alpha=0.2)

        # Add a caption to the right bottom corner of the outside of the plot
        caption = '*Vertical dashed lines indicate chromosome boundaries.'
        ax.text(1.0, -0.2, caption, transform=ax.transAxes, ha='right', va='bottom', fontsize=12)

        # Set the title of the plot
        f.suptitle('k-mer Based GWAS Manhattan Plot for ' + snakemake.params["pheno"], fontsize=title_fontsize) 

        # Set the x and y axis labels
        plt.xlabel('Chromosome', fontsize=label_fontsize) # Set the x axis label
        plt.ylabel(r"$-log_{10}{(P)}$", fontsize=label_fontsize) # Set the y axis label
        plt.xticks(fontsize = tick_fontsize)
        plt.yticks(fontsize = tick_fontsize)

        # Adjust the plot layout
        plt.tight_layout() 

      ## Saving the plot as pdf
      print("Plotting is done. Saving the plot...")
      plt.savefig(snakemake.output["manhattan_plot"], dpi=dpi)

import pysam
from optparse import OptionParser

parser = OptionParser("$prog [options]")
parser.add_option("-i", "--infile", dest="infile", help="Input GWAS+freq file", default=None, type="string")
parser.add_option("-p", "--minp", dest="minp", help="Minimum p-value allowed", default=1, type="float")
parser.add_option("-o", "--outfile", dest="outfile", help="Output file", default=None, type="string")
parser.add_option("-s", "--sep", dest="sep", help="Number of valid SNPs separating each printed SNP", default=None, type="int")
(options,args) = parser.parse_args()


infile = pysam.Tabixfile(options.infile, mode='r')
outfile = open(options.outfile,"w")

readheader = infile.header
for x in readheader:
    header = x

header = header.split("\t")
header = list(filter(lambda a: a != "REF" and a != "ALT" and a != "ANC" and a != "DER" and a!= "SE" and a != "PVAL", header))
header = "\t".join(header)

outfile.write(''.join(header) + '\n')

i=0
for line in infile.fetch():
    fields = line.strip().split("\t")
    Pval = float(fields[9])
    if Pval > options.minp:
        if i == options.sep:
            i = 0
            finalvec = fields[0:3]+[fields[7]]+fields[10:]
            outfile.write("\t".join(finalvec)+ '\n')
        i += 1

import pysam
import numpy as np
from optparse import OptionParser
import subprocess

parser = OptionParser("$prog [options]")
parser.add_option("-i", "--infile", dest="infile", help="Input GWAS+freq file", default=None, type="string")
parser.add_option("-b", "--bedfile", dest="bedfile", help="Bed file with LD partitions (def None)", default=None, type="string")
parser.add_option("-p", "--maxp", dest="maxp", help="Maximum p-value allowed", default=1, type="float")
parser.add_option("-o", "--outfile", dest="outfile", help="Output file", default=None, type="string")
(options,args) = parser.parse_args()


bedfile = open(options.bedfile,"r")
infile = pysam.Tabixfile(options.infile, mode='r')
outfile = open(options.outfile,"w")
logmaxp = np.log10(options.maxp)

readheader = infile.header
for x in readheader:
    header = x

header = header.split("\t")
header = list(filter(lambda a: a != "REF" and a != "ALT" and a != "ANC" and a != "DER" and a!= "SE" and a != "PVAL", header))
header = "\t".join(header)
outfile.write(''.join(header) + '\n')

for line in bedfile:
    regfields = line.strip("\n").split("\t")
    regchr = regfields[0].strip()
    regstart = int(regfields[1])
    regend = int(regfields[2])

    CurrLogPval = 0
    CurrPval = None
    CurrBest = None
    try:
        for elem in infile.fetch(regchr,regstart-1,regend):
            fields = elem.strip().split("\t")
            Pval = float(fields[9])
            if Pval == 0.0:
                Pval = 1e-20
            LogPval = np.log10(Pval)
            if LogPval < CurrLogPval:
                CurrBest = fields[0:3]+[fields[7]]+fields[10:]
                CurrLogPval = LogPval
                CurrPval = Pval
    except:
        continue

    if CurrLogPval < logmaxp and CurrBest != None:
        outfile.write("\t".join(CurrBest) + '\n')

import os,sys
import numpy as np
import gzip
from collections import Counter, defaultdict
import fileinput
import operator
import pickle
import time
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.pyplot import cm
import editdistance
import numpy as np
import pysam
from masq_helper_functions import tabprint
from masq_helper_functions import reverseComplement
from masq_helper_functions import convert_cigar_string
from masq_helper_functions import setup_logger, load_snv_table, write_snv_table

########################################################################
# Start timer
t0 = time.time()
# Setup log file
log = setup_logger(snakemake.log,'check_loci')
log.info('Starting process')

########################################################################
# INPUT FILES AND PARAMETERS
log.info('Getting input files and parameters from snakemake object')
# WGS path to bam
# Now allows multiple BAMs
if isinstance(snakemake.input.bam,list) and isinstance(snakemake.params.wgs_name,list):
    WGS_BAM = snakemake.input.bam
    WGS_NAME = snakemake.params.wgs_name
elif isinstance(snakemake.input.bam,str) and isinstance(snakemake.params.wgs_name,str):
    WGS_BAM = [snakemake.input.bam]
    WGS_NAME = [snakemake.params.wgs_name]
else:
    logger.error("WGS_BAM and WGS_BAM should be lists or strings")

# # Input SNV table
SNV_table = open(snakemake.input.SNV_table,'r')

########################################################################
# OUTPUT FILES
log.info('Getting output files from snakemake object')
# Output plots
list_of_plot_files = snakemake.output.plots
# Updated SNP file
updated_SNV_table = snakemake.output.new_SNV_table

# Make plot folder if it doesn't exist
plotfolder = os.path.dirname(list_of_plot_files[0])
os.makedirs(plotfolder, exist_ok=True)

########################################################################
# Plotting colors
BASE_COLORS = ["#00ABBA",
               "#9ACA3C",
               "#F26421",
               "#672D8F"]

########################################################################
## load the genome sequence from a pickle file
log.info('Loading reference genome pickle')
seq_pickle = snakemake.params.wgs_ref
seq_dic = pickle.load(open(seq_pickle, 'rb'))
log.info('Done loading reference genome pickle')

########################################################################
# Load the input SNV table
log.info('Loading SNV table')
snv_info = load_snv_table(SNV_table)
for key,value in snv_info.items():
    log.debug(key)
    log.debug(tabprint(value))
log.info('Done loading SNV table')

########################################################################

# Process input file to extend it with strand info and get coordinates of seq
log.info('Parsing specific SNV info fields')

target_info = []
target_locs_array = [list(map(int, x.split(";") ) ) if len(x)>0 else [] for x in snv_info['target_locs']]

for i,loc in enumerate(snv_info['loc']):
    chrom=snv_info['chr'][i]
    pos = int(snv_info['posi'][i])
    target_locs = target_locs_array[i]
    aml_loc = target_locs[0]
    log.debug('Position: %d' % pos)
    log.debug('AML Loc: %d' % aml_loc)
    int_seq = snv_info['trimmed-target-seq'][i]
    length = len(int_seq)

    if chrom=="0":
        strand="+"
        start=1
        end=100
        targets=target_locs
    elif ('strand' in snv_info.keys()) and ('fragment-start' in snv_info.keys()) and ('fragment-end' in snv_info.keys()) and ('add-targets' in snv_info.keys()):
        strand=snv_info['strand'][i]
        start=snv_info['fragment-start'][i]
        end=snv_info['fragment-end'][i]
        end=snv_info['add-targets'][i]
    else:
        log.info('Inferring strand, start, and end from match to reference genome')
        # identifying start and end genomic positions of the targeted region
        # and getting target sequence from ref genome
        top_start = pos - 1 - aml_loc
        top_end = top_start + length
        top_match = seq_dic[chrom][top_start:top_end]
        # if the target sequence was on the bottom strand...
        bottom_start = pos - (length - aml_loc)
        bottom_end = bottom_start + length
        bottom_match = reverseComplement(seq_dic[chrom][bottom_start:bottom_end])

        log.debug('Target seq : %s' % int_seq)
        log.debug('Top match  : %s' % top_match)
        log.debug('Btm match  : %s' % bottom_match)

        # check which strand the sequence came from and set coordinates
        if top_match == int_seq:
            log.debug('Locus %s: positive strand' % (loc))
            strand = "+"
            start = top_start
            end = top_end
            targets = target_locs
        else:
            log.debug('Locus %s: negative strand' % (loc))
            strand = "-"
            start = bottom_start
            end = bottom_end
            targets = [length - x - 1 for x in target_locs]

    target_info.append([chrom, strand, start, end, targets])

########################################################################
# Now go through WGS data and add non-ref bases to the input file...
log.info('Finding additional variant positions from WGS BAM')
BASES = ["A", "C", "G", "T"]
#sample_names = [WGS_NAME] # EDIT TO ALLOW MULTIPLE BAMS...
BASE2INT = dict([x[::-1] for x in enumerate(BASES)])
MIN_QUAL = 20
MIN_MAP = 40
buff = 30
region_buffer = 30


all_targets = []
loc_ind = 0  # counter for which loci we're on
#
# load each region from the target info read in and processed
for chrom, strand, true_start, true_end, true_targets in target_info:

    if chrom=="0": # not human / mouse sequence (GFP seq for example)
        all_targets.append([])
        image_filename = list_of_plot_files[loc_ind]
        fig = plt.figure(figsize=(20, 18))
        plt.savefig(image_filename, dpi=200, facecolor='w', edgecolor='w',
                    papertype=None, format=None,
                    transparent=False)
        plt.close()
        loc_ind += 1
    else:

        log.debug('Locus %d' % (loc_ind+1))
        image_filename = list_of_plot_files[loc_ind]
        loc_ind += 1

        # add buffer to target sequence on either side and update values
        start = true_start - region_buffer
        end = true_end + region_buffer
        targets = [x + region_buffer for x in true_targets]
        # get local sequence
        local_seq = seq_dic[chrom][start:end]
        log.debug('Chrom: %s' % chrom)
        log.debug('start: %d' % start)
        log.debug('end: %d' % end)
        log.debug('Local_seq: %s' % local_seq)
        L = len(local_seq)
        # also in integer form
        local_int = np.array([BASE2INT[x] for x in local_seq])
        # intialize for keeping track of variants at this position
        has_something = np.zeros(L, dtype=bool)

        # for each sample
        fig = plt.figure(figsize=(20, 18))
        locus_string = "%s:%d-%d %s" % (chrom, start, end, strand)
        fig.suptitle(locus_string, fontsize=20)
        ind = 1
        for sample_name,sample_bam_fn in zip(WGS_NAME,WGS_BAM):
            log.info('Plotting locus %d' % loc_ind)
            # Set up the plot
            ax1 = fig.add_subplot(len(WGS_NAME), 1, ind)
            ax2 = ax1.twinx()
            ax1.tick_params(axis='both', labelsize=15)
            ax2.tick_params(axis='both', labelsize=15)
            ind += 1

            # open the bam file
            sample_bam = pysam.AlignmentFile(sample_bam_fn, "rb")

            # for each position of interest
            # get an iterator that walks over the reads that include our event
            read_iterator = sample_bam.fetch(
                chrom, start - buff, end + buff)  # ref has chr1,chr2,etc.

            # store as pairs
            read_names = []
            read_pair_dict = dict()
            for read in read_iterator:
                try:
                    read_pair_dict[read.qname].append(read)  # if pair is there
                except:
                    read_pair_dict[read.qname] = [read]  # if pair is absent
                    read_names.append(read.qname)  # keys to dictionary

            N = len(read_names)  # number of read pairs in this region
            # store an integer at every position for each read(pair) in interval
            # matrix (number of reads (N) by number of positions (end-start))
            read_stack = np.zeros(shape=(N, end - start), dtype=int)
            # get all read pairs
            for rcounter, read_name in enumerate(read_names):
                reads = read_pair_dict[read_name]
                for read in reads:
                    read_sequence = read.query_sequence
                    read_qualities = read.query_qualities
                    # get_aligned_pairs:
                    # returns paired list of ref seq position and read seq position
                    for rindex, pindex in read.get_aligned_pairs():
                        # rindex and pindex return ALL positions, including soft-clipped
                        # rindex: position in read (1-150 for example)
                        # pindex: position in reference genome
                        # if there's a base in the read
                        # and we are in range
                        if rindex is not None and pindex is not None:
                            # Separated these lines, added pindex None check
                            if pindex >= start and pindex < end:
                                # compute the base score
                                base_score = read_qualities[rindex]
                                # if it's good enough
                                if base_score >= MIN_QUAL:
                                    # check if there is already a value in place
                                    current = read_stack[rcounter, pindex - start]
                                    base = read_sequence[rindex]
                                    if base == "N":
                                        baseint = 0
                                    else:
                                        # A-1, C-2, G-3, T-4
                                        baseint = BASE2INT[base] + 1
                                    # if there is no value, store the value
                                    if current == 0:
                                        read_stack[rcounter, pindex - start] = baseint
                                    else:
                                        # if there is a mismatch between the two reads,
                                        # set value back to 0
                                        # this value is just for the 2 paired reads
                                        if current != baseint:
                                            read_stack[rcounter, pindex - start] = 0
            summary = []
            # iterating over numpy array - by rows
            # transpose first to iterate over positions
            for x in read_stack.transpose():
                # gets counts of N,A,C,G,T per position as array
                # append to summary list
                summary.append(np.bincount(x, minlength=5))
            # convert summary to array, fills by row
            # drop the N count, and transpose again
            # .T tranposes, only if num dimensions>1
            summary = np.array(summary)[:, 1:].T
            # now we have base (4) by position array as summary
            # base_cover: coverage at each position (sum of A/C/G/T counts)
            base_cover = np.sum(summary, axis=0)
            # base_ratio: A/C/G/T counts over total coverage
            # aka frequency of each base
            base_ratio = summary.astype(float) / np.maximum(1, base_cover)
            # update has something
            # EDIT - 0 coverage is not has_something??
            has_something += ( (base_ratio[local_int, np.arange(L)] < 0.9) & (base_cover>3) ) # reference ratio is less than 0.9

            ########################################################################
            # Plot variants in each region

            # plot the coverage first
            ax1.plot(base_cover, color='k', label='coverage', alpha=0.5)
            # draw lines for boundaries of event
            # and targets
            if strand == "+":
                ax2.axvline(region_buffer-0.5, color='g', lw=2)
                ax2.axvline(L-region_buffer-0.5, color='g', linestyle='--', lw=2)
            else:
                ax2.axvline(region_buffer-0.5, color='g', linestyle='--', lw=2)
                ax2.axvline(L-region_buffer-0.5, color='g', lw=2)
            for pos in targets:
                ax1.axvline(pos, color='r', lw=2)

            # Plot colored circle for each base at position vs. frequency
            for BASE_COLOR, BASE, yvals in zip(BASE_COLORS, BASES, base_ratio):
                ax2.plot(yvals, 'o', markersize=13, label=BASE, color=BASE_COLOR)
                base_filter = local_int == BASE2INT[BASE]  # same as ref base
                x = np.arange(L)
                # label ref bases with  black circle
                ax2.plot(x[base_filter], yvals[base_filter], 'o',
                         markersize=13, mfc="None", mec='black', mew=2)
            ax1.set_ylabel(sample_name, fontsize=25)

            # for non-ref sites...
            # label number of reads for non-ref base in red
            for hs_ind in np.where(has_something)[0]:
                base_counts, total = summary[:, hs_ind], base_cover[hs_ind]
                ref_base = local_int[hs_ind]
                for bind in range(4):
                    # loop through A/C/G/T, find the non-zero counts
                    if bind != ref_base and base_counts[bind] > 0:
                        # text_string = r"$\frac{%d}{%d}$" % (
                            # base_counts[bind],total)
                        # text_string = r"%d/%d" % (base_counts[bind], total)
                        text_string = r"%d" % (base_counts[bind])
                        # add the count to the plot
                        ax2.text(hs_ind + 1, base_ratio[bind, hs_ind], text_string,
                                 fontsize=20, color='red', weight='semibold')
            ax2.set_xlim(-10, 10+L)
            ax2.set_ylim(0, 1)
            ax2.legend(loc="upper center", numpoints=1, ncol=4,
                       fontsize=18, columnspacing=0.8, handletextpad=-0.2)
            ax1.legend(loc="lower center", numpoints=2,
                       fontsize=18, columnspacing=0.5)

        plt.tight_layout(rect=(0, 0, 1, 0.98))
        plt.savefig(image_filename, dpi=200, facecolor='w', edgecolor='w',
                    papertype=None, format=None,
                    transparent=False)
        plt.close()

    ########################################################################
        # Add non-ref het/homo sites to "target" list
        rlen = true_end - true_start
        new_targets = np.where(has_something)[0] - region_buffer
        if strand == "-":
            new_targets = true_end - true_start - new_targets - 1

        # new_targets>=0 (in original target region)
        new_targets = new_targets[(new_targets >= 0) * (new_targets < rlen)]
        all_targets.append(new_targets)

    ########################################################################

# Add new het/homo non-ref sites to original input file
log.info('Writing updated SNV info table')
outfile = open(updated_SNV_table, 'w')

snv_info['add-targets']=[]
snv_info['strand']=[]
snv_info['fragment-start']=[]
snv_info['fragment-end']=[]
print("Start loop")
for i, (new_targets, more_info) in enumerate(zip(all_targets, target_info)):
    print(i)
    prev_targets = list(map(int, snv_info['target_locs'][i].split(";")))
    add_targets = [x for x in new_targets if x not in prev_targets]
    snv_info['add-targets'].append(";".join(list(map(str, add_targets))))
    snv_info['strand'].append(more_info[1])
    snv_info['fragment-start'].append(more_info[2])
    snv_info['fragment-end'].append(more_info[3])
print("End loop")
print(snv_info)

write_snv_table(snv_info,outfile)

########################################################################
# Close files
outfile.close()

########################################################################
# End timer
t1 = time.time()
td = (t1 - t0) / 60
log.info("Done in %0.2f minutes" % td)