import os, subprocess, time
import tables as tb

time_steps = 500;
datasets = 17;
dataset_size = 50000;

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)


def create_file(filename):
    f = tb.openFile(filename, "w")
    arr = tb.numpy.empty(dataset_size, dtype='float64')  # buffer
    for tstep in range(time_steps):
        gname = "Timestep_"+str(tstep)
        g = f.createGroup(f.root, gname)
        for data_id in range(datasets):
            d = f.createArray(g, "dataset_"+str(data_id), arr)
    f.close()


def test_file(filename):
    f = tb.openFile(filename, "r")
    for tstep in range(time_steps):
        gname = "/Timestep_"+str(tstep)
        g = f.getNode(gname)
        for data_id in range(datasets):
            d = f.getNode(g, "dataset_"+str(data_id))
            data = d[:]
    f.close()

if __name__ == '__main__':
    import getopt, sys
    opts, pargs = getopt.getopt(sys.argv[1:], 's', ['create-file'])
    filename = pargs[0]
    show = False
    create = False
    for option in opts:
        if option[0] == '--create-file':
            create = True
        elif option[0] == '-s':
            show = True

    tref = time.time()
    if create:
        print "Creating file...", filename
        create_file(filename)
        if show:
            show_stats("File written", tref)
    print "Reading file...", filename
    for i in range(10):
        test_file(filename)
        if show:
            show_stats("Read iteration #%d" % i, tref)