#!/usr/bin/env python

helptext = '''This script will retrieve paralog nucleotide (CDS) sequences for a specified
gene in all samples located in namelist.txt. It writes all the (unaligned) sequences to stdout.
If a sample does not have paralogs for that gene, the sequence in the FNA directory is retrieved instead.'''

import os,sys,argparse
from Bio import SeqIO

def retrieve_seqs(path,name,gene):
	seqs_to_write = None
	if os.path.isdir(os.path.join(path,name,gene,name,'paralogs')):
		seqs_to_write = [x for x in SeqIO.parse(os.path.join(path,name,gene,name,'paralogs','{}_paralogs.fasta'.format(gene)),'fasta')]
		num_seqs = str(len(seqs_to_write))
	elif os.path.isfile(os.path.join(path,name,gene,name,'sequences','FNA','{}.FNA'.format(gene))):
		seqs_to_write = SeqIO.read(os.path.join(path,name,gene,name,'sequences','FNA','{}.FNA'.format(gene)),'fasta')
		num_seqs = "1"
	
	if seqs_to_write:
		SeqIO.write(seqs_to_write,sys.stdout,'fasta')
	else:
		num_seqs = "0"
	
	return num_seqs
			

def main():
	parser = argparse.ArgumentParser(description=helptext,formatter_class=argparse.RawTextHelpFormatter)
	parser.add_argument('namelist',help='Text file containing list of HybPiper output directories, one per line.')
	parser.add_argument('gene',help="Name of gene to extract paralogs")
	
	args = parser.parse_args()
	
	namelist = [x.rstrip() for x in open(args.namelist)]
	num_seqs = []
	for name in namelist:
                path, name = os.path.split(name)
                if not name:
                        path,name = os.path.split(path)
	        num_seqs.append(retrieve_seqs(path,name,args.gene))
	sys.stderr.write("{}\t{}\n".format(args.gene,"\t".join(num_seqs)))
	
if __name__ == "__main__":main()