"""Testbed for open/close PyTables files.

This uses the HotShot profiler.

"""

from __future__ import print_function
import os
import sys
import getopt
import pstats
import cProfile as prof
import time
import subprocess  # From Python 2.4 on
import tables

filename = None
niter = 1


def show_stats(explain, tref):
    "Show the used memory"
    # Build the command to obtain memory info (only for Linux 2.6.x)
    cmd = "cat /proc/%s/status" % os.getpid()
    sout = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).stdout
    for line in sout:
        if line.startswith("VmSize:"):
            vmsize = int(line.split()[1])
        elif line.startswith("VmRSS:"):
            vmrss = int(line.split()[1])
        elif line.startswith("VmData:"):
            vmdata = int(line.split()[1])
        elif line.startswith("VmStk:"):
            vmstk = int(line.split()[1])
        elif line.startswith("VmExe:"):
            vmexe = int(line.split()[1])
        elif line.startswith("VmLib:"):
            vmlib = int(line.split()[1])
    sout.close()
    print("WallClock time:", time.time() - tref)
    print("Memory usage: ******* %s *******" % explain)
    print("VmSize: %7s kB\tVmRSS: %7s kB" % (vmsize, vmrss))
    print("VmData: %7s kB\tVmStk: %7s kB" % (vmdata, vmstk))
    print("VmExe:  %7s kB\tVmLib: %7s kB" % (vmexe, vmlib))


def check_open_close():
    for i in range(niter):
        print(
            "------------------ open_close #%s -------------------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        fileh.close()
        show_stats("After closing file", tref)


def check_only_open():
    for i in range(niter):
        print("------------------ only_open #%s -------------------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        show_stats("Before closing file", tref)
        fileh.close()


def check_full_browse():
    for i in range(niter):
        print("------------------ full_browse #%s -----------------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        for node in fileh:
            pass
        fileh.close()
        show_stats("After full browse", tref)


def check_partial_browse():
    for i in range(niter):
        print("------------------ partial_browse #%s --------------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        for node in fileh.root.ngroup0.ngroup1:
            pass
        fileh.close()
        show_stats("After closing file", tref)


def check_full_browse_attrs():
    for i in range(niter):
        print("------------------ full_browse_attrs #%s -----------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        for node in fileh:
            # Access to an attribute
            klass = node._v_attrs.CLASS
        fileh.close()
        show_stats("After full browse", tref)


def check_partial_browse_attrs():
    for i in range(niter):
        print("------------------ partial_browse_attrs #%s --------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        for node in fileh.root.ngroup0.ngroup1:
            # Access to an attribute
            klass = node._v_attrs.CLASS
        fileh.close()
        show_stats("After closing file", tref)


def check_open_group():
    for i in range(niter):
        print("------------------ open_group #%s ------------------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        group = fileh.root.ngroup0.ngroup1
        # Access to an attribute
        klass = group._v_attrs.CLASS
        fileh.close()
        show_stats("After closing file", tref)


def check_open_leaf():
    for i in range(niter):
        print("------------------ open_leaf #%s -----------------------" % i)
        tref = time.time()
        fileh = tables.open_file(filename)
        leaf = fileh.root.ngroup0.ngroup1.array9
        # Access to an attribute
        klass = leaf._v_attrs.CLASS
        fileh.close()
        show_stats("After closing file", tref)


if __name__ == '__main__':

    usage = """usage: %s [-v] [-p] [-n niter] [-O] [-o] [-B] [-b] [-g] [-l] [-A] [-a] [-E] [-S] datafile
              -v verbose  (total dump of profiling)
              -p do profiling
              -n number of iterations for reading
              -O Check open_close
              -o Check only_open
              -B Check full browse
              -b Check partial browse
              -A Check full browse and reading one attr each node
              -a Check partial browse and reading one attr each node
              -g Check open nested group
              -l Check open nested leaf
              -E Check everything
              -S Check everything as subprocess
              \n""" % sys.argv[0]

    try:
        opts, pargs = getopt.getopt(sys.argv[1:], 'vpn:OoBbAaglESs')
    except:
        sys.stderr.write(usage)
        sys.exit(0)

    progname = sys.argv[0]
    args = sys.argv[1:]

    # if we pass too much parameters, abort
    if len(pargs) != 1:
        sys.stderr.write(usage)
        sys.exit(0)

    # default options
    verbose = 0
    silent = 0  # if silent, does not print the final statistics
    profile = 0
    all_checks = 0
    all_system_checks = 0
    func = []

    # Checking options
    options = ['-O', '-o', '-B', '-b', '-A', '-a', '-g', '-l']

    # Dict to map options to checking functions
    option2func = {
        '-O': 'check_open_close',
        '-o': 'check_only_open',
        '-B': 'check_full_browse',
        '-b': 'check_partial_browse',
        '-A': 'check_full_browse_attrs',
        '-a': 'check_partial_browse_attrs',
        '-g': 'check_open_group',
        '-l': 'check_open_leaf',
    }

    # Get the options
    for option in opts:
        if option[0] == '-v':
            verbose = 1
        elif option[0] == '-p':
            profile = 1
        elif option[0] in option2func:
            func.append(option2func[option[0]])
        elif option[0] == '-E':
            all_checks = 1
            for opt in options:
                func.append(option2func[opt])
        elif option[0] == '-S':
            all_system_checks = 1
        elif option[0] == '-s':
            silent = 1
        elif option[0] == '-n':
            niter = int(option[1])

    filename = pargs[0]

    tref = time.time()
    if all_system_checks:
        args.remove('-S')  # We don't want -S in the options list again
        for opt in options:
            opts = "%s \-s %s %s" % (progname, opt, " ".join(args))
            # print "opts-->", opts
            os.system("python2.4 %s" % opts)
    else:
        if profile:
            for ifunc in func:
                prof.run(ifunc + '()', ifunc + '.prof')
                stats = pstats.Stats(ifunc + '.prof')
                stats.strip_dirs()
                stats.sort_stats('time', 'calls')
                if verbose:
                    stats.print_stats()
                else:
                    stats.print_stats(20)
        else:
            for ifunc in func:
                eval(ifunc + '()')

    if not silent:
        print("------------------ End of run -------------------------")
        show_stats("Final statistics (after closing everything)", tref)
