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

if '-h' in sys.argv:
	print 'Reads list generated by getlocus.py and KEEP only the sequences from the given alignment of the proteins in the list\n\n\tSintaxe: python keeplocus.py file.list alignment.fa'
	sys.exit()

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

list_locus = []
for line in listfile:
	if line.count('-') > 1:
		cut_point = line[:line.find('-')+1].find('-') + line.find('-')
		list_locus.append(line[:cut_point])
		
	else:
		list_locus.append(line.replace('\n',''))

print list_locus
seq_dic, seq_list = bitk.fastareader(sys.argv[2])

found_list = []

output = ''
for seq in seq_list:
	#seq = seq.replace('/','-')
	locus = seq.split('-')[1]
	#organism = seq.split('-')[0] + '-' + seq.split('-')[1]
	if locus in list_locus:
		output += '>' + seq + '\n' + seq_dic[seq] + '\n'
		found_list.append(locus)

for locus in list_locus:
	if locus not in found_list:
		print "Not found: " + locus
	

datafile = open(sys.argv[2][:-3] + '.new.fa', 'w')
datafile.write(output)
datafile.close()
listfile.close()
		
