<commit_before>#!/usr/bin/env python
# Built-in python packages
from __future__ import division # This modifies Python 2.7 so that any expression of the form int/int returns a float 
import getopt
import sys
import os
# Project-specific packages
from queries_with_ref_bases import query_contains_ref_bases

def get_accession_map(fasta_path):
	'''
	Suppose in the FASTA file used as reference for makeblastdb, there are n sequences.
	Magic-BLAST labels these sequences as the order in which they appear when outputting in tabulated format.
	For example, suppose that the SNP rs0001 appears second from the top in the FASTA file.
	Then Magic-BLAST assigns the label '1' to rs0001 whenever it appears in an alignment.
	This function returns a dictionary that serves as a map from integers (in string datatype) to accessions
	e.g. accession_map['1'] == rs0001 in our above example.
	Inputs
	- (str) fasta_path: path to the FASTA file used as reference for makeblastdb
	Outputs
	- (dict) accession_map: the map from integers to accessions
	'''
        accession_map = {}
        with open(fasta_path,'r') as fasta:
                id_number = 0
                for line in fasta:
                        if line[0] == ">":
                                accession = line[1:].rstrip()
                                accession_map[str(id_number)] = accession
                                id_number += 1
        return accession_map


def get_mbo_paths(directory):
	'''
	Given a directory, retrieves the paths of all files within the directory whose extension is .mbo
	Inputs
	- (str) directory: directory to be scoured for mbo files
	Outputs
	- paths: a dictionary where the keys are accessions and the values are paths
	'''
	paths = {}
	for file in os.listdir(directory):
		if file.endswith(".mbo"):
			# We only want the accession number, not the extension as well
			accession = os.path.basename(file).split('.')[0]
			path = os.path.join(directory,file)
			paths[accession] = path
	return paths

def get_sra_alignments(paths,accession_map):
	'''
	Given a list of paths as described in the function get_mbo_paths, retrieves the BTOP string for each
	alignment.
	Inputs
	- paths: a list of pairs where the first entry of the pair is the accession and the second is the path
	- (dict) accession_map: the map between integers and accessions
	Outputs
	- a dictionary where keys are SRA accessions and the values are alignment dictionaries
	'''
	sra_alignments = {}
	for accession in paths:
		path = paths[accession]
		alignments = []	
		with open(path,'r') as mbo:
			for line in mbo:
				tokens = line.split()
				# Skip the line if it is commented, the number of fields isn't equal to 25 or
				# the query read was not aligned
				if line[0] != "#" and len(tokens) == 25 and tokens[1] != "-":
					var_acc = accession_map[ tokens[1] ]
					ref_start = int(tokens[8])
					ref_stop = int(tokens[9])
					if ref_start > ref_stop:
						temp = ref_start
						ref_start = ref_stop
						ref_stop = temp
					btop = tokens[16]
					alignment = { 'var_acc': var_acc, 'ref_start': ref_start,\
						      'ref_stop': ref_stop, 'btop': btop }
					alignments.append( alignment )
		sra_alignments[accession] = alignments
	return sra_alignments

def get_var_info(path):
	'''
	Retrieves the flanking sequence lengths for the SNP sequences			
	Inputs
	- path: path to the file that contains the flanking sequence lengths
	Outputs
	- var_info: a dictionary where the keys are SNP accessions and the values are tuples where the first entry\
                    is the start position of the variant, the second is the stop position of the variant and \
                    the third is the length of the variant sequence
	'''
	var_info = {}
	with open(path,'r') as input_file:
		for line in input_file:
			tokens = line.split()
			if len(tokens) == 4:
				accession = tokens[0]
				start = int(tokens[1])
				stop = int(tokens[2])
				length = int(tokens[3])
				var_info[accession] = {'start':start,'stop':stop,'length':length}
	return var_info

def call_variants(var_freq):
	'''
	Determines which variants exist in a given SRA dataset given the number of reads that do and do not contain
	the var
	Inputs
	- var_freq: a dict where the keys are SNP accessions and the values are dicts which contain the frequency of
		    reads that do and reads that do not contain the SNP
	Outputs
	- variants: a list which contains the SNP accessions of those SNPs which exist in the SRA dataset
	'''
	variants = []
	for var_acc in var_freq:
		frequencies = var_freq[var_acc]
		true = frequencies['true']
		false = frequencies['false']
		try:
			percentage = true/(true+false)
			if percentage > 0.4: # For now, we use this simple heuristic.
				variants.append(var_acc)
		except ZeroDivisionError: # We ignore division errors because they correspond to no mapped reads
			pass
	return variants

def get_sra_variants(sra_alignments,var_info):
	'''
	For all SRA accession, determines which variants exist in the SRA dataset	
	Inputs
	- sra_alignments: dict where the keys are SRA accessions and the values are lists of alignment dicts
	- var_info: dict where the keys are variant accessions and the values are information concerning the variants
	Outputs
	- variants: dict where the keys are the SRA accessions and the values are lists which contain the accessions
		    of variants which exist in the SRA datasets
	'''
	variants = {}
	for sra_acc in sra_alignments:
		alignments = sra_alignments[sra_acc]
		var_freq = {}
		for alignment in alignments:
			var_acc = alignment['var_acc']
			# Get the flank information
			info = var_info[var_acc]
			if var_acc not in var_freq:
				var_freq[var_acc] = {'true':0,'false':0}
			# Determine whether the variant exists in the particular SRA dataset
			var_called = query_contains_ref_bases(alignment,info)
			if var_called == True:
				var_freq[var_acc]['true'] += 1	
			elif var_called == False:
				var_freq[var_acc]['false'] += 1	
		sra_variants = call_variants(var_freq) 
		variants[sra_acc] = sra_variants	
	return variants

def create_tsv(variants,output_path):
	'''
	Creates a TSV file containing the set of variants each SRA dataset contains.
	Inputs
	- variants: dict where the keys are the SRA accessions and the values are lists which contain the accessions
		    of variants which exist in the SRA datasets
	- output_path: path to where to construct the output file
	'''
	with open(output_path,'w') as tsv:
		header = "SRA	Variants\n"
		tsv.write(header)
		for sra_acc in variants:
			line = "%s" % (sra_acc)		
			sra_variants = variants[sra_acc]
			for var_acc in sra_variants:
				line = "%s	%s" % (line, var_acc)
			tsv.write(line)
			tsv.write('\n')

if __name__ == "__main__":
	help_message = "Description: Given a directory with Magic-BLAST output files where each output file\n" \
                     + "             contains the alignment between an SRA dataset and known variants in a human\n" \
                     + "             genome, this script determines which variants each SRA dataset contains\n" \
                     + "             using a heuristic."<commit_msg>modified call_variants.py to differentiate between homozygous and heterozygous snps<commit_after>#!/usr/bin/env python
# Built-in python packages
from __future__ import division # This modifies Python 2.7 so that any expression of the form int/int returns a float 
import getopt
import sys
import os
# Project-specific packages
from queries_with_ref_bases import query_contains_ref_bases

def get_accession_map(fasta_path):
	'''
	Suppose in the FASTA file used as reference for makeblastdb, there are n sequences.
	Magic-BLAST labels these sequences as the order in which they appear when outputting in tabulated format.
	For example, suppose that the SNP rs0001 appears second from the top in the FASTA file.
	Then Magic-BLAST assigns the label '1' to rs0001 whenever it appears in an alignment.
	This function returns a dictionary that serves as a map from integers (in string datatype) to accessions
	e.g. accession_map['1'] == rs0001 in our above example.
	Inputs
	- (str) fasta_path: path to the FASTA file used as reference for makeblastdb
	Outputs
	- (dict) accession_map: the map from integers to accessions
	'''
        accession_map = {}
        with open(fasta_path,'r') as fasta:
                id_number = 0
                for line in fasta:
                        if line[0] == ">":
                                accession = line[1:].rstrip()
                                accession_map[str(id_number)] = accession
                                id_number += 1
        return accession_map


def get_mbo_paths(directory):
	'''
	Given a directory, retrieves the paths of all files within the directory whose extension is .mbo
	Inputs
	- (str) directory: directory to be scoured for mbo files
	Outputs
	- paths: a dictionary where the keys are accessions and the values are paths
	'''
	paths = {}
	for file in os.listdir(directory):
		if file.endswith(".mbo"):
			# We only want the accession number, not the extension as well
			accession = os.path.basename(file).split('.')[0]
			path = os.path.join(directory,file)
			paths[accession] = path
	return paths

def get_sra_alignments(paths,accession_map):
	'''
	Given a list of paths as described in the function get_mbo_paths, retrieves the BTOP string for each
	alignment.
	Inputs
	- paths: a list of pairs where the first entry of the pair is the accession and the second is the path
	- (dict) accession_map: the map between integers and accessions
	Outputs
	- a dictionary where keys are SRA accessions and the values are alignment dictionaries
	'''
	sra_alignments = {}
	for accession in paths:
		path = paths[accession]
		alignments = []	
		with open(path,'r') as mbo:
			for line in mbo:
				tokens = line.split()
				# Skip the line if it is commented, the number of fields isn't equal to 25 or
				# the query read was not aligned
				if line[0] != "#" and len(tokens) == 25 and tokens[1] != "-":
					var_acc = accession_map[ tokens[1] ]
					ref_start = int(tokens[8])
					ref_stop = int(tokens[9])
					if ref_start > ref_stop:
						temp = ref_start
						ref_start = ref_stop
						ref_stop = temp
					btop = tokens[16]
					alignment = { 'var_acc': var_acc, 'ref_start': ref_start,\
						      'ref_stop': ref_stop, 'btop': btop }
					alignments.append( alignment )
		sra_alignments[accession] = alignments
	return sra_alignments

def get_var_info(path):
	'''
	Retrieves the flanking sequence lengths for the SNP sequences			
	Inputs
	- path: path to the file that contains the flanking sequence lengths
	Outputs
	- var_info: a dictionary where the keys are SNP accessions and the values are tuples where the first entry\
                    is the start position of the variant, the second is the stop position of the variant and \
                    the third is the length of the variant sequence
	'''
	var_info = {}
	with open(path,'r') as input_file:
		for line in input_file:
			tokens = line.split()
			if len(tokens) == 4:
				accession = tokens[0]
				start = int(tokens[1])
				stop = int(tokens[2])
				length = int(tokens[3])
				var_info[accession] = {'start':start,'stop':stop,'length':length}
	return var_info

def call_variants(var_freq):
	'''
	Determines which variants exist in a given SRA dataset given the number of reads that do and do not contain
	the var
	Inputs
	- var_freq: a dict where the keys are SNP accessions and the values are dicts which contain the frequency of
		    reads that do and reads that do not contain the SNP
	Outputs
	- variants: a list which contains the SNP accessions of those SNPs which exist in the SRA dataset
	'''
	variants = {'heterozygous':[],'homozygous':[]}
	for var_acc in var_freq:
		frequencies = var_freq[var_acc]
		true = frequencies['true']
		false = frequencies['false']
		try:
			percentage = true/(true+false)
			if percentage > 0.8: # For now, we use this simple heuristic.
				variants['homozygous'].append(var_acc)
			elif percentage > 0.3:
				variants['heterozygous'].append(var_acc)
		except ZeroDivisionError: # We ignore division errors because they correspond to no mapped reads
			pass
	return variants

def get_sra_variants(sra_alignments,var_info):
	'''
	For all SRA accession, determines which variants exist in the SRA dataset	
	Inputs
	- sra_alignments: dict where the keys are SRA accessions and the values are lists of alignment dicts
	- var_info: dict where the keys are variant accessions and the values are information concerning the variants
	Outputs
	- variants: dict where the keys are the SRA accessions and the values are lists which contain the accessions
		    of variants which exist in the SRA datasets
	'''
	variants = {}
	for sra_acc in sra_alignments:
		alignments = sra_alignments[sra_acc]
		var_freq = {}
		for alignment in alignments:
			var_acc = alignment['var_acc']
			# Get the flank information
			info = var_info[var_acc]
			if var_acc not in var_freq:
				var_freq[var_acc] = {'true':0,'false':0}
			# Determine whether the variant exists in the particular SRA dataset
			var_called = query_contains_ref_bases(alignment,info)
			if var_called == True:
				var_freq[var_acc]['true'] += 1	
			elif var_called == False:
				var_freq[var_acc]['false'] += 1	
		sra_variants = call_variants(var_freq) 
		variants[sra_acc] = sra_variants	
	return variants

def create_tsv(variants,output_path):
	'''
	Creates a TSV file containing the set of variants each SRA dataset contains.
	Inputs
	- variants: dict where the keys are the SRA accessions and the values are lists which contain the accessions
		    of variants which exist in the SRA datasets
	- output_path: path to where to construct the output file
	'''
	with open(output_path,'w') as tsv:
		header = "SRA\tHeterozygous SNPs\tHomozygous SNPs\n"
		tsv.write(header)
		for sra_acc in variants:
			line = "%s" % (sra_acc)		
			sra_variants = variants[sra_acc]
			line = line + "\t"
			for var_acc in sra_variants['heterozygous']:
				line = line + var_acc + ","
			line = line + "\t"
			for var_acc in sra_variants['homozygous']:
				line = line + var_acc + ","
			tsv.write( line.rstrip() )
			tsv.write('\n')

if __name__ == "__main__":
	help_message = "Description: Given a directory with Magic-BLAST output files where each output file\n" \
                     + "             contains the alignment between an SRA dataset and known variants in a human\n" \
                     + "             genome, this script determines which variants each SRA dataset contains\n" \
                     + "             using a heuristic."
{'fraction_non_alphanumeric': 0.07709175632152401, 'fraction_numerical': 0.006030288494483656, 'mean_word_length': 3.2424418604651164, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1508705', 'n_tokens_mistral': 4871, 'n_tokens_neox': 4306, 'n_words': 1695}