#! /usr/bin/env python

'''
Created on Aug 8, 2013

@author: smirarab
'''
import sys
from pasta.alignment import CompactAlignment
import argparse
from copy import copy


if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(description='Manipulate Alignments') 
    parser.add_argument('-infile', nargs='?', type=argparse.FileType('r'), default=sys.stdin,
                         help='name of the input file (default: standard input)')
    parser.add_argument('-outfile', nargs='?', type=argparse.FileType('w'), default=sys.stdout,
                         help='name of the output file (default: standard output)')
    parser.add_argument('-informat', nargs='?', default="FASTA", choices = ["COMPACT3", "FASTA"],
                         help='format of the input file (default: FASTA)')
    parser.add_argument('-outformat', nargs='?',  default="FASTA", choices = ["COMPACT3", "FASTA", "PHYLIP"],
                         help='format of the output file (default: FASTA)')    
    parser.add_argument('-masksites', metavar='N', type=int, 
                         help='sites with less than N non-gap characters will be masked out') 
    parser.add_argument('-filterfragments', metavar='N', type=int, 
                         help='sequences with less than N non-gap sequences will be removed')        
    parser.add_argument('-rename', metavar='MappingFile', type=argparse.FileType('r'), 
                         help='Rename sequences, according to the mapping file generated by PASTA')      
#     parser.add_argument('actions', nargs='*', 
#                          help='a list of actions (separated by space). Supported actions are: mask (see maskmin), rename (see namemap)')
  
    args = parser.parse_args()
    alg = CompactAlignment()
    alg.read_file_object(args.infile,args.informat)
    
#     if args.actions:
#         print >>sys.stderr, "will perform the following actions: %s" %",".join(args.actions)
#     else:
#         print >>sys.stderr, "No actions will be taken. Input in %s will be written to output in %s." %(args.informat,args.outformat)
    
    if args.masksites:
        alg.mask_gapy_sites(args.masksites)
        
    if args.filterfragments:
        rem = []
        for k, v in alg.items():
            if len(v.seq) < args.filterfragments:
                rem.append(k)
        for k in rem:
            alg.pop(k, None)
        
    if args.rename:
        lines=[x for x in args.rename.readlines() if x!="\n"]
        namemap=dict(list(zip((x.strip() for i,x in enumerate(lines) if i%2 == 0),
                         (x.strip() for i,x in enumerate(lines) if i%2 == 1))))
        names = copy(list(alg.keys()))
        for k in names:
            alg[namemap[k]] = alg.pop(k)
            
    alg.write(args.outfile, args.outformat)
    
