#! /usr/bin/env python
# ______________________________________________________________________
#
# Filename: chunkstats.py
#
# Created: 17-JUL-2001,   Harry Melanson
#
# ______________________________________________________________________

"""Prints statistics related to chunks in a DSPACK file.

USAGE:

  chunkstats [-hsn] dspackfile

OPTIONS:

  -h, --help     --  Print help text and exit.
  -n nevents     --  Specify number of events to process
                     [default = 1]
  -s             --  Sort by size [default = sort by name]

Prints a summary of the average size of chunks found
in a DSPACK file.  Note that the sizes are non-compressed
values.  Compression will make the actual files smaller.

By default the first event is examined.  If the '-n nevents'
option is used, the average size of the chunks for the specified
number of events is calculated.

The results are by default sorted by chunk name.  To sort by
chunk size, use the '-s' option.

"""

# ______________________________________________________________________

import sys
import glob
import string
import d0om

# ______________________________________________________________________

__author__ = 'Harry Melanson, melanson@fnal.gov'
__version__ = '01.00.00'
__filename__ = 'chunkstats.py'
__doc__ = __doc__ % vars()

# ______________________________________________________________________

class ChunkStat:

    def __init__(self, name):

	self.name = name
	self.size = 0.0
	self.nchunks = 0.0

    def update(self, size):

	self.size = self.size + size
	self.nchunks = self.nchunks + 1.0

    def avgsize(self):
	if self.nchunks != 0:
	    return self.size / self.nchunks
	else:
	    return 0.0

    def __repr__(self):
	s = []
	s.append('%30s' % self.name)
	s.append(", avg size: " + ('%10.1f' % self.avgsize()) + " bytes")
	return string.join(s,"")

# ______________________________________________________________________
#

def sortByName(a,b):
    if a.name < b.name:
	return -1
    elif a.name > b.name:
	return 1
    else:
	return 0

def sortBySize(a,b):
    if a.avgsize() < b.avgsize():
	return -1
    elif a.avgsize() > b.avgsize():
	return 1
    else:
	return 0

# ______________________________________________________________________

def process_command():
    """Process the command."""

    # ______________________________________________________________________
    #
    #  Initialize
    # ______________________________________________________________________
    #

    from scriptutil import getOptions
    optlist, args = getOptions('sn:', [], __doc__)

    #  Process command line arguments
    
    if len(args) == 0:
	print "\nERROR: No DSPACK file specified."
        print "\nNAME: ", (os.path.basename(sys.argv[0]))
        print  __doc__
	sys.exit(-1)

    dsfile = args[0]
    if len(glob.glob(dsfile)) == 0:
	print "\nERROR: DSPACK file", dsfile, "does not exist.\n"
	sys.exit(-1)
    elif len(glob.glob(dsfile)) != 1:
	print "\nERROR: More than one file found?\n"
	sys.exit(-1)

    #  Process command options

    nevents = 1
    sortoption = 'byname'

    for opt in optlist:

	
	if opt.flag == '-n':
	    nevents = string.atoi(opt.value)
	elif opt.flag == '-s':
	    sortoption = 'bysize'

    # ______________________________________________________________________
    #

    print '\n\nProcessing', nevents, 'events - please wait...'

    # ______________________________________________________________________
    #
    #  Process the DSPACK file, keeping track of chunk statistics
    # ______________________________________________________________________
    #

    stats = {}
    total_size = 0.0

    #  Open file

    file = d0om.d0Stream(dsfile)

    #  Loop over requested number of events

    for i in range(0,nevents):

	event = file.read()
	if (event == None):
	    break

	chunks = event['_event']['_mapByType']

	#  Loop over chunks

	for key in chunks.keys():
	    if not stats.has_key(key):
		stats[key] = ChunkStat(key)

	    #  Keep track of the size of each chunk
	    for chunk in chunks[key]:
		stats[key].update(chunk.size())
		total_size = total_size + chunk.size()

	print 'Event', i+1


    # ______________________________________________________________________
    #
    #  Print summary
    # ______________________________________________________________________
    #

    chunklist = []
    for key in stats.keys():
	chunklist.append(stats[key])

    if (sortoption == 'byname'):
	chunklist.sort(sortByName)
    else:
	chunklist.sort(sortBySize)

    for chunk in chunklist:
	print chunk

    print "\nAvg event size: ", total_size / nevents, 'bytes'


#
#  Self-test

if __name__ == '__main__':
    process_command()
    sys.exit(0)