#!/usr/bin/env python

import sys
import argparse
from Bio import SeqIO

def main(argv, wayout):
	if not len(argv):
		argv.append('-h')
	parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__)
	parser.add_argument('-a','--alignment', help="alignment in fasta format")
	parser.add_argument('-s','--sites', help="comma separate list of sites, like '75,79,150'")
	parser.add_argument('-i','--keep-index', action="store_true", help="sites given are already indexed for python (meaning minus 1)")
	args = parser.parse_args(argv)
    
    sites = [int(s) for s in args.sites.split(",")]
    if not args.keep_index:
        sites = [i-1 for i in sites]

    for seqrec in SeqIO.parse(args.alignment, "fasta"):
        seqsites = [seqrec.seq[s] for s in sites]
        posstring = "-".join(seqsites)
        seqrec.id = "{}_{}".format( str(seqrec.id), posstring )
        wayout.write( seqrec.format("fasta") )

if __name__ == "__main__":
	main(sys.argv[1:], sys.stdout)