#!/usr/bin/env python

'''
================================================================================
###  Motif-Mark ###
A program to convert sequences motifs to a visual representation of the locations of motifs in sequences.  This script can take DNA or RNA sequence input.  It requires
a fasta file and a text file with the motifs.  This script can take up to 13 motifs. This script returns a visual image displaying the sequence
 as a line with exons overlapped on the line as empty black boxes and motifs as solid colored areas.  It alo returns a text list of motif locations on the genes given.
 For each sequence the entire fasta header line (with chromosome, location and such) is given.
 ===============================================================================

================================================================================
 To start the program in an Ubuntu shell:
 Make sure that you are in the directory where your files are at. Ex: cd  /mnt/c/Users/david/OneDrive/Documents/Winter_Genomics_2020
conda install -c conda-forge pycairo
Start program, including entering arg parse for fasta file and motifs.text file as follows:
./motif_mark.py -s -s <directory path>/<sequence_file>.fasta -m <directory path>/<motif_file.text
For me that, once in the correct directory, is ./motif_mark.py -s ./motif_mrker_sequence.fasta -m ./motif_marker_motifs.txt.text
Your motif file should look something like the following:

AACTGYat
YYYSYYN
AAaaggcR
ccgtgtca
================================================================================
'''
#===============================================================================
#Import necessary libraries
#===============================================================================
import cairo
import re
import argparse
import sys
import math
import numpy as np
import matplotlib.cm as cm
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
from re import finditer

#===============================================================================_
#  Stand-alone Functions
#===============================================================================

def getargs():
	#A function to input sequence file, motif sequence,  returns these as arguments, then variable, uses argparse module.
	parser = argparse.ArgumentParser(description = 'input fasta file and motif list .txt file and output an image of where these motifs are along particular genes')
	parser.add_argument('-s', '--seqfile', required =True, type= str, help= "Enter -s or --seqfile followed by your absolute fasta file path")
	parser.add_argument('-m', '--motifs', required = True, type = str, help= "-m 0r --motifs followed by the absolute file path to a text file with one motif per line.")

	return parser.parse_args()

# Set variables for Arg items
args = getargs()
seqfile = args.seqfile
motifs = args.motifs

#unit test
#for line in open(seqfile,'r'):
#    print(line)
#for line in open(motifs,'r'):
#    print(line)
#===============================================================================

def motif_iupac_translater(motifs_list):
	#A Function to translate a IUPAC Motif lists into amino acids.  This is set to handle  DNA motifs, with lower or upper case [exons] and
	#lower or upper case IUPAC variable nucleotide symbols. This function turns a DNA motif-list, upper or lower case into a list of IUPAC
	#equivalents in regex ready form that can be compared to a DNA or an RNA sequence'''
	search_motifs =[]
#	translation_dict = {"A":"[Aa]", "a":"[Aa]", "}
	for string in motifs_list:
		string = string.replace("A", "a")
		string = string.replace("a","[Aa]")
		string = string.replace("G","g")
		string = string.replace("g","[Gg]")
		string = string.replace("C","c")
		string = string.replace("c","[Cc]")
		string = string.replace("T","t")
		string = string.replace("t","[TtUu]")
		string = string.replace("R","r")
		string = string.replace("r","[AaGg]")
		string = string.replace("Y","y")
		string = string.replace("y","[CcTtUu]")
		string = string.replace("S","s")
		string = string.replace("s","[GgCc]")
		string = string.replace("W","w")
		string = string.replace("w","[TtUu]")
		string = string.replace("K","k")
		string = string.replace("k","[GgTtUu]")
		string = string.replace("S","s")
		string = string.replace("s","[GgCc]")
		string = string.replace("W","w]")
		string = string.replace("w","[AaTtUu]")
		string = string.replace("M","m")
		string = string.replace("m","[AaCc]")
		string = string.replace("B","b")
		string = string.replace("b","[GgCcTtUu]")
		string = string.replace("D","d")
		string = string.replace("d","[AaGgTtUu]")
		string = string.replace("H","h")
		string = string.replace("h","[AaCcTtUu]")
		string = string.replace("V","v")
		string = string.replace("v","[AaCcGg]")
		string = string.replace("N","n")
		string = string.replace("n","[GgCcAaTtUu]")
		regexmotif =  string
		search_motifs.append(regexmotif)

	return(search_motifs)
#Unit Tests
#search_motifs = motif_iupac_translater(motifs_list)
#print(search_motifs)
#===============================================================================_

def gene_capture(seqfile):
	#A function to capture genes from a fasta file as a dictionary with gene_id as key and sequence as value.  This uses the seqfile from arg parse
	with open(seqfile,'r')as fh:
		for line in fh:
			line = line.strip()
			if line.startswith(">"): #captures IDs
				seq_name = line[1:]
				genes_dict[seq_name] = []
			else:
				sequence = line
				#print(sequence)
				#print(genes)
				genes_dict[seq_name].append(sequence) #captures sequences

		for key in genes_dict:#turns value (sequence lines in a list) into value(one string of sequences)
			new_value = "".join(genes_dict[key])
			genes_dict[key] = new_value #reassigns value
	return(genes_dict)
	seqfile.close()


#===============================================================================
# Get a list of motifs from the inputted motifs_list.txt
#===============================================================================
motifs_list = [line.rstrip('\n') for line in open(motifs)]

#===============================================================================
# create a debugger object to identify prints for debugging
#===============================================================================
debugger = print

#===============================================================================
# Create a list of regex ready motifs and a dictionary that pairs the iupac motifs with the regex ready versions
#===============================================================================
# This calls on the motif_iupac_translater finction
search_motifs = motif_iupac_translater(motifs_list)
#debugger(search_motifs)
motif_regex_dict = dict(zip(motifs_list,search_motifs))

#===============================================================================
# Call gene_capture function
#===============================================================================
genes_dict = {}
genes = gene_capture(seqfile)

#===============================================================================
# Measure the lengths of the sequences and get the longest sequence
#===============================================================================
#This will be used to calibrate the pycairo drawing sheet.

values_list = list(genes.values())
sequence_lengths = []
for item in values_list:
#	debugger(item)
	sequence_length = len(item)
	#debugger(sequence_length)
	sequence_lengths.append(sequence_length)
longest_sequence = max(sequence_lengths)
#debugger(longest_sequence)

#===============================================================================
# Make color dictionary to match each motif with a color
#===============================================================================
color_list = [(1,0,0),(0,0.8,0),(0,0,1),(0.4,0,0.8),(1,0.2,1),(1,1,0.2),(0.4,1,1),(1,0.5,0),(0.3,0.6,0),(0.6,0,0.3),(1,0.6,0.6),(0.4,0.2,0),(0.5,0.5,0.5)]
color_dict = {}
for i in range(len(motifs_list)):
	color_dict[motifs_list[i]] = [color_list[i]]
#debugger(color_dict)

#===============================================================================
# Set other global variables, while outside the loops
#===============================================================================
motif_locations = {}  #To store the location in the sequence, where motifs are found, paired with the motif
for header in genes:  #popuate motif_locations with fasta file header
	motif_locations[header] = []
	#debugger(motif_locations)
exon_locations = {}  #To store the locations of exons
margin = 100  # Sets margin for the cairo page
spans_info = []

#===============================================================================
#  Set Up Cairo and Title for motif marker Page
#===============================================================================
surface = cairo.SVGSurface("Motif_marker_Results.svg",longest_sequence + margin, 100 * len(genes) + 200)   #sets the drawing surface by the longest_sequence and the number of genes in the fasta file.
context = cairo.Context(surface)
context.move_to(margin - 50, 30)
context.set_line_width(1)  #This and the next four commands create the Title
context.set_font_size(20)
context.select_font_face("Calibri")
context.text_path("Motifs found in these sequences. (Exons are marked with a rectangle.)")
context.stroke()


#===============================================================================
#Loops to get name of genes(from header), exons and motifs in each gene and draw them
#===============================================================================
spacer = 70  # A counter to put spaces between sequences
for header in genes:  #Do the following for each gene sequence in the fasta file.
	sequence_length = len(genes[header])  #get the length of the sequences to be used to calibrate images to the space
	#debugger(sequence_length)
	#  These next three steps draw a line for each gene, the length of the sequence.
	context.set_source_rgb(0,0,0)
	context.move_to(margin, spacer)  #Set up sequence lines
	thin_line = context.set_line_width(1)
	context.line_to(margin + sequence_length, spacer)   #The x is determined by the sequence length, the y is spaced from the last sequence
	context.stroke()
	#  These next six steps use the entire header lines of the fasta file to name the genes
	context.move_to(margin, spacer - 10)
	context.set_line_width(1)
	context.set_font_size(10)
	context.select_font_face("calibri")
	context.text_path(header)
	context.stroke()
    # get the exon locations and map them
	exon_locations = []  # initialize local storage of exon locations for one gene to mark exons
	string = genes[header]  # make it easy to place sequence from genes dictionary into regeax
	#debugger(string)
	exon = re.search(r"[ACTGU]+", string)  # Find Exons, which will be capitalized
	#debugger(exon)
	if exon:
		exon_locations.append(exon.span())  # find span of exons
	#debugger(exon_locations)
	# This nested loop draws box for each exon found in exon_locations
	for item in exon_locations:
		context.set_line_width(1)
		context.set_source_rgb(0,0,0)
		context.rectangle(item[0] + margin,spacer -5,item[1] - item[0], 6)
		context.stroke()
	# This nested loop does the following for each motif within each gene
	for motif in motif_regex_dict:
		motif_locations[motif] = [] #sets key of motif_locations to motif in list
		#debugger(motif_locations)
		search_motif = str(motif_regex_dict[motif])  #sets regex ready motif as string
		#debugger(search_motif)
		#debugger(type(search_motif))
		spans = []  # set local list to collect motif location span
		matches = re.finditer(search_motif, string, re.IGNORECASE)  #Finds motifs
		#debugger(matches)
        #debugger(type(matches))
		context.set_source_rgb(*color_dict[motif][0])
		for match in matches:  #For each motif found do the following
			span = match.span()  #gets span of motif from finditer iterable search objest
			#debugger(span)
			group = match.group()  #may not use. Saved in case of need
			#print(group)
			spans.append(span)  #Add span to local list
			#debugger(spans)
			#debugger(spans_info)
			#  This nested loop draws a thick vertical looking colored line for each span
			context.set_line_width(10)
			context.move_to(span[0] + margin, spacer)
			context.line_to(span[1] + margin, spacer)
			context.stroke()
		spans_info.append((header,motif,spans))
		#debugger(spans_info)
	spacer += 100  #Adds to the spacer for readability

#===============================================================================
# Draw key from color_dict, displaying motif and the associated color
#===============================================================================

for i in range(len(color_dict)):
    #Draw motif name.  Ex: YYYYYY
    context.set_line_width(1)
	context.set_source_rgb(0,0,0)
	context.move_to(margin, spacer)
	context.set_font_size(10)
	context.select_font_face("Calibri")
	context.text_path(list(color_dict.keys())[i] + " = ")
	context.stroke()
    #Draw color block to match motif
	context.set_line_width(15)
	context.set_source_rgb(*list(color_dict.values())[i][0])
	context.move_to(margin + 80, spacer)
	context.line_to(margin + 100, spacer)
	context.stroke()
	spacer += 20
surface.finish()

#===============================================================================
# Create and output text list of motif locations within each gene.
#===============================================================================

f = open("Motifs_text.txt", "w")
spans_text_info = ""
for item in spans_info:
	spans_text_info += (str(item) + "\n")
f.write("List of motifs found for each gene.  \n" + spans_text_info)
