#!/usr/bin/env python

"""
Add all (potentially gigantic) histograms in a group of files.
"""

import dashi
import tables
import os, sys, operator, shutil

from optparse import OptionParser
parser = OptionParser(usage="%prog [OPTIONS] infiles outfile", description=__doc__)
parser.add_option("--blocksize", dest="blocksize", type=int, default=2048)

opts, args = parser.parse_args()
if len(args) < 2:
	parser.error("You must specify at least one output and one input file")
infiles, outfile = args[:-1], args[-1]

if os.path.exists(outfile):
	parser.error("%s already exists!" % outfile)

shutil.copy(infiles[0], outfile)

from collections import defaultdict
paths = defaultdict(list)
for fname in infiles[1:]:
	with tables.openFile(fname) as hdf:
		for group in hdf.walkNodes(where='/', classname='Group'):
			if 'ndim' in group._v_attrs: # a dashi histogram
				path = group._v_pathname
				paths[path].append(fname)

def histadd(sourceGroup, destGroup, blocksize=1):
	"""
	Add dashi histograms stored in HDF5 groups
	
	:param blocksize: operate on blocksize I/O chunks at a time
	"""
	for arr in '_h_bincontent', '_h_squaredweights':
		source = sourceGroup._v_children[arr]
		dest = destGroup._v_children[arr]
		chunksize = blocksize*reduce(operator.mul, dest.chunkshape)
		size = reduce(operator.mul, dest.shape)
		for i in range(0, size, chunksize):
			dest[i:i+chunksize] += source[i:i+chunksize]
	for prop in 'nentries', 'nans', 'nans_wgt', 'nans_sqwgt':
		destGroup._v_attrs[prop] += sourceGroup._v_attrs[prop]

with tables.openFile(outfile, 'a') as ohdf:
	for path, fnames in paths.iteritems():
		print(path)
		destGroup = ohdf.getNode(path)
		for fname in fnames:
			with tables.openFile(fname) as hdf:
				histadd(hdf.getNode(path), destGroup, opts.blocksize)

[((12, 9, 12, 83), 'optparse.OptionParser', 'OptionParser', (), '', False, 'from optparse import OptionParser\n'), ((20, 3, 20, 26), 'os.path.exists', 'os.path.exists', ({(20, 18, 20, 25): 'outfile'}, {}), '(outfile)', False, 'import os, sys, operator, shutil\n'), ((23, 0, 23, 32), 'shutil.copy', 'shutil.copy', ({(23, 12, 23, 22): 'infiles[0]', (23, 24, 23, 31): 'outfile'}, {}), '(infiles[0], outfile)', False, 'import os, sys, operator, shutil\n'), ((26, 8, 26, 25), 'collections.defaultdict', 'defaultdict', ({(26, 20, 26, 24): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((50, 5, 50, 34), 'tables.openFile', 'tables.openFile', ({(50, 21, 50, 28): 'outfile', (50, 30, 50, 33): '"""a"""'}, {}), "(outfile, 'a')", False, 'import tables\n'), ((28, 6, 28, 28), 'tables.openFile', 'tables.openFile', ({(28, 22, 28, 27): 'fname'}, {}), '(fname)', False, 'import tables\n'), ((55, 8, 55, 30), 'tables.openFile', 'tables.openFile', ({(55, 24, 55, 29): 'fname'}, {}), '(fname)', False, 'import tables\n')]