#!/usr/bin/env python

import sys
import os


if not len(sys.argv) == 4:
    print("Usage: splitFasta <fasta file> <number of parts> <output prefix>")
    print("\nThis script divides a fasta file into several fasta files with approx equal size.")
    sys.exit(2)


fastaFile = sys.argv[1] # Input fasta file
N = int(sys.argv[2]) # number of parts to divide the file into
prefix = sys.argv[3] # output files prefix "prefix_0001.fa"


size = os.path.getsize(fastaFile)

i = 1 # counter for current part
posIn = 0 # input file position (total number of bytes read)

fout = open(prefix+"_%04d.fa" % (i,), 'w')

with open(fastaFile, 'r') as fin:
  for line in fin:
    # Check if new sequence starts and current position exceeds 
    # the expected end position of the current part
    if( line.startswith('>') and (posIn > size*i/N) ):
      # start with the next part
      fout.close()
      i += 1
      fout = open(prefix+"_%04d.fa" % (i,), 'w')
    
    fout.write(line)
    posIn += len(line)

fout.close()
