#! python

"""Split the standard input into a bunch of output files of
approximately specfied size.  (You specify size, not the number
of output files.)

Usage: split_list.py filesize pathname

The pathname must contain "%d" somewhere.  That will be replaced
with an integer from 1 to n.

Splitting is done randomly.
"""

import os
import sys
import math
import heapq
from gmisclib import die
from gmisclib import threaded_io


class szlist(object):
	def __init__(self):
		self.sz = 0.0
		self.contents = []

	def add(self, sz, x):
		self.sz += sz
		self.contents.append((sz, x))

	def dump(self):
		return [x for (sz, x) in self.contents]

	def __cmp__(self, other):
		tmp = cmp(len(self.contents), len(other.contents))
		if tmp:
			return tmp
		return cmp(self.sz, other.sz)


def do_split(fnroot, write_pattern, lab_pattern, np=1):
	NMAX = 500
	assert np < NMAX
	assert '%d' in fnroot
	assert '%s' in write_pattern
	assert '%s' in lab_pattern

	def getstat(fn):
		return os.stat(lab_pattern % fn).st_size

	fnames = [q.strip() for q in sys.stdin]
	if len(fnames) <= np:
		n = len(fnames)
	else:
		n = np * (1+int(math.sqrt(len(fnames)/float(np))))
	if n >= NMAX:
		n = NMAX-1
	assert n > 0
	hq = [ szlist() for i in range(n)]
	heapq.heapify(hq)
	for (fn, siz) in threaded_io.pairmap(getstat, fnames):
		if isinstance(siz, threaded_io.CaughtException):
			die.warn("Fail on %s: %s" % (fn, siz))
			continue
		tmp = heapq.heappop(hq)
		esiz = siz**2
		tmp.add(esiz, fn)
		heapq.heappush(hq, tmp)

	tp = threaded_io.Thread_poolW(10)
	for (i,olist) in enumerate(hq):
		tp.start(_writelines, args=(i, olist, write_pattern, fnroot))
	tp.join()


def _writelines(i, olist, write_pattern, fnroot):
	fd = open(fnroot % i, 'w')
	for q in olist.dump():
		fd.write(write_pattern % q + '\n')


if __name__ == '__main__':
	arglist = sys.argv[1:]
	if arglist[0] == '-np':
		np = int(arglist[1])
		arglist = arglist[2:]
		assert np > 0
	else:
		np = 1
	do_split(arglist[0], arglist[1], arglist[2], np=np)
