#!/usr/bin/python

import sys, random, codecs, os.path, os


def main(source_wbst, out_dir, folds, seed):
	print "WBST Cross-validator"
	if not os.path.exists(out_dir):
		os.makedirs(out_dir)
	dest = out_dir

	folds = int(folds)

	wbst = []

	source = codecs.open(source_wbst,"r",'utf8')

	header = ""

	pairs = 0
	for line in source:
		if pairs>0:
			wbst.append(line)
		else:
			header = line.strip()
		pairs = pairs + 1
	source.close()

	seedval = int(seed)
	random.seed(seedval)
	random.shuffle(wbst)

	n = int(pairs/folds)
	i = 0
	while (i<folds):
		print "FOLD " + str(i) + " OF " + str(folds)
		path = os.path.join(dest, str(i))
		if not os.path.exists(path):
			os.makedirs(path)
		teach = codecs.open(os.path.join(path,"teach.csv"),"w",'utf8')
		test = codecs.open(os.path.join(path,"test.csv"), "w", 'utf8')
		teach.write(header+'\n')
		test.write(header+'\n')
		j = 0
		while (j<len(wbst)):
			if (j < i*n or j> (i+1) * n):
				test.write(wbst[j])
			else:
				teach.write(wbst[j])
			j = j + 1
		i = i + 1
		teach.close()
		test.close()
		
if __name__ == "__main__":
	import sys, argparse
	parser = argparse.ArgumentParser(description='WBST Cross-validator')
	#positional arguments
	parser.add_argument('source_wbst',  action='store', 
		help='Path to WBST test')
	parser.add_argument('out_dir',  action='store', 
		help='Target directory')
	parser.add_argument('folds',  action='store', 
		help='Number of folds')
	parser.add_argument('seed',  action='store', 
		help='Random seed')
	args = parser.parse_args()
	
	main(args.source_wbst, args.out_dir, args.folds, args.seed)
