# -*- coding: utf-8 -*-
"""
Splits single wikipedia txt file onto multiples smaller files.

SplitWiki by Bartosz Broda, bartosz.broda@gmail.com
"""
import codecs, os
import argparse

	
def main():
	parser = argparse.ArgumentParser(description='Splits (Westbury) Wikipedia txt dump onto singleton files.')
	parser.add_argument('-w', '--wiki', dest = 'wikifile', required=True, 
			help = 'Wikipedia txt file to process')
	parser.add_argument('-d', '--dirprefix', dest = 'dirprefix', default = 'outwiki/wiki_', 
			required=True, help = 'Prefix for output directory, default: outwiki/wiki_')
	parser.add_argument('-f', '--filesprefix', dest = 'fileprefix', default = 'xx', 
			help = 'Prefix for file names, default: xx')
	parser.add_argument('-m', '--maxfiles', dest = 'max_files', default = 10000, type=int,
			help = 'Maximum number of files in one directory, default: 10000')
	args = parser.parse_args()

	file_no = 0
	dir_no = 0

	file_in = args.wikifile
	dirprefix = args.dirprefix
	fileprefix = args.fileprefix
	max_files = args.max_files

	curdir = (dirprefix + str(dir_no) )
	os.makedirs(curdir)

	out = codecs.open( os.path.join(curdir, fileprefix + '%.7d' % file_no), 'w', 'utf8')

	no = 0
	for line in codecs.open(file_in, 'r', 'utf8'):
		if line == '---END.OF.DOCUMENT---\n':
			no += 1
			out.close()
			file_no += 1
			if no % max_files == 0:
				dir_no += 1
				print "Dir %d (doc: %d)" % (dir_no, no)
				curdir = (dirprefix + str(dir_no) )
				os.makedirs(curdir)
			out = codecs.open( os.path.join(curdir, fileprefix + '%.7d' % file_no), 'w', 'utf8')
			if no % 1000 == 0:
				print "Processed ", no / 1000, 'k documents'
		else:
			out.write(line)

	if not out.closed:
		out.close()

if __name__ == '__main__':
	main()
