##################################
#                                #
# Last modified 2018/09/22       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string

def run():

    if len(sys.argv) < 3:
        print 'usage: python %s inputfilename <max piece size> <outfile prefix>' % sys.argv[0]
        sys.exit(1)

    inputfilename = sys.argv[1]
    MaxSeq = int(sys.argv[2])
    outputfileprefix = sys.argv[3]

    lineslist = open(inputfilename)

    sequenceDict={}
    i=0
    ID=''
    sequence=''
    for line in lineslist:
        if i % 100000 == 0:
            print i, 'lines processed'
        i+=1
        if line[0] == '>':
            if ID != '':
                sequenceDict[ID] = sequence
            ID = line.strip()
            sequence=''
        else:
            sequence = sequence+line.strip()
    sequenceDict[ID] = sequence

    print 'finished inputting sequence'
   
    IDList = sequenceDict.keys()
    IDList.sort()
    TotalSequence = 0
    k = 1
    for i in range(0,len(IDList),MaxSeq):
        print 'outputting file', k
        outfile = open(outputfileprefix + '.' + str(k) + '.fa', 'w')
        for j in range(i,min(len(IDList),i + MaxSeq)):
            outfile.write(IDList[j] + '\n')
            outfile.write(sequenceDict[IDList[j]] + '\n')
        k+=1
        outfile.close()

run()