##################################
#                                #
# Last modified 04/22/2013       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import os

def run():

    if len(sys.argv) < 2:
        print 'usage: python %s fasta output_directory [-fields 0,1,...,N separator]' % sys.argv[0]
        sys.exit(1)

    fasta = sys.argv[1]
    outdir = sys.argv[2]

    doSplit = False
    if '-fields' in sys.argv:
        doSplit = True
        SplitFields = []
        fields = sys.argv[sys.argv.index('-fields') + 1].split(',')
        for f in fields:
            SplitFields.append(int(f))
        separator = sys.argv[sys.argv.index('-fields') + 2]

    cmd = 'mkdir ' + outdir

    try:
        os.system(cmd)
    except:
        pass

    outfile = ''
    linelist = open(fasta)
    for line in linelist:
        if line.startswith('>'):
            chr = line.strip().split('>')[1]
            if doSplit:
                fields = chr.split(separator)
                chr = ''
                for f in SplitFields:
                    chr = chr + fields[f] + separator
                chr = chr[0:-len(separator)]
            if outfile == '':
                pass
            else:
                outfile.close()
            outfile = open(outdir + '/' + chr + '.fa','w')
            outfile.write(line)
        else:
            outfile.write(line)

    outfile.close()

run()