#!/usr/local/bin/python
import sys
import os.path
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
import sys
import os.path
import csv
import matplotlib
import matplotlib.pyplot as plt
#compares all sequences in aligned fasta file to first sequence in set
#written by M.Cotten WTSI and EMC
if len(sys.argv)!= 3:
	print("Usage: python HiLiter.py alignment_table panel_size")
	sys.exit()

print "Running HiLiter.py really fast"

alignment_table= sys.argv[1]
outprefix = os.path.splitext(alignment_table)[0]
panel_size= sys.argv[2]
#prepare list of alignments to plot
alignment_list=[]
starting_csv = open(alignment_table, "rU")
reader1 = csv.reader(starting_csv)
for row in reader1:
	alignment_list.append(row[0])
fig = plt.figure()
plot_index=1
for alignment in alignment_list:
	seq_names =[]
	sequences =[]
	for record in SeqIO.parse(open(alignment, "rU"), "fasta"):
		seq_names.append(record.id)
		sequence_string = str(record.seq)
		sequences.append(sequence_string)
	total_number_sequences = len(sequences)
	genome_length = len(sequences[0])

	if total_number_sequences >18 and total_number_sequences <=21:
		ytick_label_size = 4
	else:
		ytick_label_size = 6
	graph_names=[]
	graph_names.append(seq_names[0])
	width_unit = 15
	gap_unit=1
	marker_size = 45
	graph_pad = 5 * marker_size

	#Generate lists of differences between ref and test sequences
	#function for comparing sequence a to b
	def vergleichen(seqA, seqB, index):
		gap_end_marker = 500000 #need value that would never be encountered.
		gap_start_position = 0
		current_gap_width = 0
		current_gap_marker = 0
		gap_dictionary = {}
		for i in range (len(seqA)):
			position = i
			if seqB[i]!=seqA[i] and seqB[i] == "A":
				color = "orangered"
				line = index
				diff_list.append((position,width_unit,line,color))
			elif seqB[i]!=seqA[i] and seqB[i] == "T":
				color = "crimson"
				line = index
				diff_list.append((position,width_unit,line,color))
			elif seqB[i]!=seqA[i] and seqB[i] == "G":
				color = "indigo"
				line = index
				diff_list.append((position,width_unit,line,color))
			elif seqB[i]!=seqA[i] and seqB[i] == "C":
				color = "slateblue"
				line = index
				diff_list.append((position,width_unit,line,color))
			elif seqB[i] == "N": #currently this maps each nt gap.
				if position-1 != gap_end_marker: #if the previous position is not gap then this is start of new gap
					gap_start_position = position
					gap_end_marker = position
					current_gap_width = gap_unit
					print "position"
					print position
					print "start of a gap"
					color = "#e6e6e6"
					line = index
					print "line"
					print line
					gap_dictionary[gap_start_position]=(current_gap_width,line,color) #move to
				else: #then it must be == to gap end marker and is an extension of a gaps
					current_gap_width = current_gap_width+gap_unit #extend current gap	by 1 unit
					gap_end_marker = position
					print "extend a gap"
					print "gap_start_position"
					print gap_start_position
					print "current_gap_width"
					print current_gap_width

					color = "#e6e6e6"
					line = index
					print "line"
					print line
					gap_dictionary[gap_start_position]=(current_gap_width,line,color) #move to gap_dictionary

		for k,v in gap_dictionary.iteritems():
			gap_start_position = k
			current_gap_width = v[0]
			line = v[1]
			color = v[2]
			print "dictionary dump"
			print gap_start_position,current_gap_width,line,color
			diff_list.append((gap_start_position,current_gap_width,line,color))

	diff_list=[]
	#call vergleichen on each pair
	for i in range(1, len(sequences)):
		vergleichen(sequences[0], sequences[i], i)

	#generate graph of differences.
	ticks=[]
	for x in range(len(sequences)):
		ticks.append(x)
	if panel_size == "normal":
		ax2 = fig.add_subplot(4,2, int(plot_index))
	elif panel_size == "large":
		ax2 = fig.add_subplot(2,2, int(plot_index))	
	ax2.set_ylim(0.9, len(sequences))
	ax2.set_xlim(-200, (genome_length+200))
	ax2.set_yticks(ticks)
	ax2.set_yticklabels(seq_names, va='bottom', size = ytick_label_size)
	ax2.tick_params(axis='x', which='major', labelsize=6)
	ax2.set_xlabel(seq_names[0]+' Genome Position', fontsize=8)
	for i in range(0, len(diff_list)):
		color = diff_list[i][3]
		ax2.broken_barh([(int(diff_list[i][0]), (diff_list[i][1]))], (int(diff_list[i][2]), 0.85), facecolors=(color),edgecolors=(color))
	plot_index = plot_index+1

plt.tight_layout()
plt.savefig(outprefix+'_differences.pdf')
print "That's All Folks!"
