#! /usr/bin/env python 
###################################
##    Davi Ortega                ##
###################################
import sys

if '-h' in sys.argv:
	print "Make the tag as the following: XX.YYY-LOCUS, XX Specie and YYY subspecie \n Sintax: pyhton seqtagtest.py datafile.dat (genome_id\tname\tlocus\tsequence\tcomplete(optional)\ttaxonomic_group(optional)\taseq_id(optional)\tgene_id(optional)\tposition(optional)\trank2(optional))\n\tFlags: \n\t-c\tonly complete genomes\n\t-t taxonomic_group \t Selects only sequences from the taxonomic_group\n\t-g only exports a list of sequence_tag;gene_id;position\n"
	sys.exit()

datafile = open(sys.argv[1],'r')

seqtag_dic = {}
#seqtag[XX.YYY-locus] = [NAME, locus]
seq_dic = {}

if '-t' in sys.argv:
	taxo_group = sys.argv[sys.argv.index('-t')+1]

for line in datafile:
	fields = line.split('\t')
	#makeing the tag:
	genome_id = fields[0].replace(' ','')
#update for Anabaena 'Nostoc azollae' 0708 ######################
        #full_name = fields[1]
        full_name = fields[1].replace("'",'')
#################################################################
	full_locus = fields[2]
	try:
		complete = int(fields[4])
	except ValueError:
		pass 
	if ('NULL' not in full_locus) and ( 'name' not in fields[1]) and ( '-c' in sys.argv and complete == 1 or '-c' not in sys.argv) and ( '-t' in sys.argv and fields[5] == taxo_group or '-t' not in sys.argv):
#update for Anabaena 'Nostoc azollae' 0708 ######################
#               name = fields[1].split(' ')
                name = fields[1].replace("'","").split(' ')
#################################################################
		while '' in name:
			name.remove('')
		while 'sp.' in name:
			name.remove('sp.')
		if len(name) < 2 or name[1].isupper() or ('-' in name[1]): 
			tag = name[0][:6]
		else:
			tag = name[0][:2] + '.' + name[1][:3]
		locus = fields[2].replace(' ','')
		tag += '.' + genome_id + '-' + locus
		#print tag
		if tag not in seqtag_dic.keys():
			seqtag_dic[tag] = [ full_name , full_locus]
			if '-g' not in sys.argv:
				seq_dic[tag] = fields[3].replace(' ','')
			else:
				seq_dic[tag] = fields[7].replace(' ','') + ';' + fields[8].replace(' ','')
		elif full_name != seqtag_dic[tag][0] or full_locus != seqtag_dic[tag][1]:
			print "Found a duplicate"
			print "Tag: " + tag
			print "Organism names: " + full_name + '\t' + seqtag_dic[tag][0]
			print "Locus names: " + full_locus + '\t' + seqtag_dic[tag][1]
			datafile.close()
			sys.exit()
		else:
			print "Something weird happened: " + tag
			print full_name
			print full_locus
			print genome_id
			print seqtag_dic[tag]
			print "\n\nPay attention to this entry:"
			print line	
	elif 'NULL' in full_locus:
		print line
					
		

datafile.close()

filename = sys.argv[1][:-4]


if '-c' in sys.argv:
	filename += '_only_complete'
if '-t' in sys.argv:
	filename += '_' + taxo_group
if '-g' not in sys.argv:
	filename += '.fa'
else:
	filename += '.gene.list'
outfile = open(filename,'w')

output = ''
for tag, seq in seq_dic.iteritems():
	if '-g' not in sys.argv:
		output += '>' + tag + '\n' + seq + '\n'
	else:
		output += tag + ';' + seq + '\n'

outfile.write(output)
outfile.close()



#for tag, list in seqtag_dic.iteritems():
#	print tag + '\t' + str(list)

print "There are " + str(len(seqtag_dic.keys())) + " proteins"
print "Done"

	

