#!/usr/bin/env
# -*- coding: utf-8 -*-
######################################################################################
####
#### Description: This script retrieve a fragment from each sequence in an aligned 
####			  fasta file (file.aligned.fa) based on the fragment provided in the 
####			  segment.fa file. (e.g. v1v3, v3v4, etc) of a reference.
####			  Note: The reference has to be the first sequence in the aligned 
####			  aligned.fa file
####
#### Usage:       python retrieveSegment.py aligned.fa segment.fa out.fa
####               
#### Author:      Yanmei Huang
#### Version:     1.0
#### Date:        2017-6-30
####
######################################################################################

__author__ = "Yanmei Huang"
__version__= "1.0"
__date__= "June 30 2017"

import sys
import os
import argparse

###############################################################################
# a function that constructs a list containing the corresponding positions of 
# each nucleotide of a reference in the alignment consensus.

def getPosInConsensus (input_fa):
	
	IN = open( input_fa, 'r')
	REF = IN.readline()
	REF = IN.readline()
	L = list()
	for i in range(0,(len(REF)-1)):
		if not REF[i] == "-" :
			L.append(i)
	return (L)
	
################################################################################
# a function that slices out segments based on an example in segment.fa
# and the corresponding indexes of its head and tail in the alignment 
# consensus

def segment_slicer( input_fa, segment_fa, output_fa, index_list):
	
	# obtain coordinates
	IN = open( input_fa, 'r' )
	SEG = open (segment_fa, 'r')
	
	segName = SEG.readline()[1:].strip()
	seg = SEG.readline()
	seg_head = seg[0:20].upper()
	seg_tail = seg[len(seg)-20: len(seg)].upper()
	seg_tail = seg_tail.strip()
	
	ref = IN.readline()
	ref = IN.readline().upper()
	ref = ref.replace("-", "")
	start = index_list[ref.find(seg_head)]
	end = index_list[ref.find(seg_tail) + 20]
	
	print (segName + "\t" + seg_head + "\t" + seg_tail + "\t" + str(start) + "\t" + str(end))

	# slice out segments of sequences from aligned fasta	
	FA = open( input_fa, 'r')
	OUT = open ( output_fa, 'w')
	
	while 1:
		header = FA.readline()
		if not header: break
		seq = FA.readline().upper()
		if not seq: break
		
		seq = seq[start : end]
		seq = seq.replace("-", "")
		
		OUT.write(header)
		OUT.write(seq + "\n")
	

############################################################################
# Main
if __name__ == "__main__":
	
	########################################################################
	# set up command line argument parsing
	
	parser = argparse.ArgumentParser(description='# This script extracts the amplicon fragment from a fasta file based on coordinates obtained from emboss primersearch output.' )
	
	parser.add_argument('in_file', action="store", help='input aligned fasta file.', type=str)
	parser.add_argument('seg_file', action="store", help='file containing an example segment', type=str)
	parser.add_argument('out_file', action="store", help='output fasta file', type=str)

	parser.add_argument('--version', action='version', version='%(prog)s 1.0')

	arguments = parser.parse_args()

	########################################################################
	# get coordinates in alignment consensus
	index_list = getPosInConsensus (arguments.in_file)
	# slice out segments
	
	segment_slicer ( arguments.in_file, arguments.seg_file, arguments.out_file, index_list )