#!/usr/bin/env python
# h5frames by David Dumas <daviddumas@gmail.com>

"""h5frames: frontend to h5topng to ease the pain of creating animations"""

import os, sys, re, string
import getopt

def usage():
    print('''
    usage: h5frames [options] <hdf5file>

    options include:
        -c <file> or --colormap <file>
           Use colormap from <file>; default is "custom.colormap" if
           it exists, or "bluered" if it does not.

        -r <m> or --range <m>
           Set the conversion range to [-<m>,<m>], or [0,<m>] if "--nozero"
           was specified.  If <m> is negative, then each frame has its range
           chosen individually by "h5topng".

        --nozero
            Do not center the range on zero.
    
        -o <name> or --output <name>
            Specifies the base name of the output images; default is
            the same as the name of the first component of the hdf5
            path.  If this option is given, then the name is instead
            the concatenation of <name> and a zero-padded zero-based
            frame number.
            
        -p <path> or --path <path>
            Specifies the path within the hdf5 file.  h5frames expects
            a path structure like the following:

            /<path>/frame01/A/
            /<path>/frame01/B/
            /<path>/frame02/A/
            /<path>/frame02/B/
            /<path>/frame03/A/
            /<path>/frame03/B/

            i.e. the first component of the path after <path> should be the "name"
            of the frame, so that the lexicographical order coincides with the
            intended order in the animation.  Any substructure should be identical
            for each frame, and the particular dataset can be selected with "-d" or
            "--dataset"

        -d <dataset> or --dataset <dataset>
            Specifies which dataset to convert within each frame.  This can include
            a hdf5 path prefix.  The full hdf5 paths converted are those of the form:

            /<path>/*/<dataset>

        -O or --overwrite
            Overwrite existing output files; default is to avoid overwriting.

        -h or --help
            Print this usage message.
    ''')

def main():
    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], "ho:d:p:r:c:OTv", ["help", "output=", "dataset=","path=","range=","colormap=","overwrite","test","verbose","nozero"])
    except getopt.GetoptError:
        # print help information and exit:
        usage()
        sys.exit(2)
    output = None
    dataset = None
    path = None
    colormap = None
    nozero = False
    h5range = 200.0
    overwrite = False
    dryrun = False
    verbose = False
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        if o in ("-o", "--output"):
            output = a
        if o in ("-p", "--path"):
            path = a
        if o in ("-d", "--dataset"):
            dataset = a
        if o in ("-r", "--range"):
            h5range = float(a)
        if o in ("-c", "--colormap"):
            colormap = a
        if o == "--nozero":
            nozero = True
        if o in ("-O", "--overwrite"):
            overwrite = True
        if o in ("-T", "--test"):
            dryrun = True
        if o in ("-v", "--verbose"):
            verbose = True

    if len(args) == 0:
        sys.stderr.write('Error: missing hdf5 filename.\n')
        usage()
        sys.exit()

    filename = args[0]

    optstring = ''
    
    if not colormap:
        if os.path.exists('custom.colormap'):
            colormap = 'custom.colormap'
        else:
            colormap = 'bluered'
    optstring += ' -c %s' % colormap

    if h5range > 0:
        if nozero:
            optstring += ' -m 0 -M %f' % (h5range)
        else:
            optstring += ' -Z -m -%f -M %f' % (h5range,h5range)
    else:
        if not nozero:
            optstring += ' -Z'

    f = os.popen("h5ls -r " + filename)
    tree = [x.split()[0].split('/')[1:] for x in f.readlines() if x.split()[1] == 'Dataset']
    f.close()

    if path:
        if path[0] == '/':
            path = path[1:]
        if path[-1] == '/':
            path = path[:-1]
        spath = path.split('/')
        findex = len(spath)
        subtree = [ x for x in tree if x[:findex] == spath ]
    else:
        findex = 0
        subtree = [ x for x in tree ]

    if dataset:
        if dataset[0] == '/':
            dataset = dataset[1:]
        if dataset[-1] == '/':
            dataset = dataset[:-1]
        sdataset = dataset.split('/')
        sstree = [ x for x in subtree if x[1:] == sdataset ]
    else:
        sstree = subtree
 
    i=0
    for p in sstree:
        if output:
            outname = output + ('%03d' % i) + '.png'
        else:
            outname = p[findex] + '.png'
        execstr = 'h5topng %s -o %s -d %s %s' % (optstring,outname,string.join(p,'/'),filename)
        if dryrun:
            print(execstr)
        else:
            if overwrite or (not os.path.exists(outname)):
                if verbose:
                    print('/%s/ --> %s' % (string.join(p,'/'),outname))
                os.system(execstr)
            else:
                if verbose:
                    print('/%s/ --> %s  [not overwriting]' % (string.join(p,'/'),outname))
                    
        i = i+1

if __name__ == '__main__':
    main()