#!/usr/bin/env python

'''
Print a tree view of a NeXus file.
Abbreviate arrays.
Replace repeated structures by a link.
'''

# (C) Joachim Wuttke, Forschungszentrum Jülich GmbH, 2014.
# Released under the GNU General Public License version 3.0 or higher.

import argparse
import numpy as np
import h5py

INDENT = "    "

def samestruct(g,h):
    '''Return true if groups g,h have the same structure.'''
    # different attributes?
    if set(g.attrs.keys()) ^ set(h.attrs.keys()):
        return False
    # different entries?
    if set(g.keys()) ^ set(h.keys()):
        return False
    # recursion over subgroups:
    for gk,gv in g.items():
        hv = h[gk]
        if isinstance( gv, h5py.Group ):
            if not samestruct( gv, hv ):
                return False
        elif isinstance( gv, h5py.Dataset ):
            if not len(gv.shape) == len(hv.shape):
                return False
    return True

def shapestr(obj):
    '''Return compact string describing the shape of given object.'''
    rank = len(obj.shape)
    out = []
    for d in range(rank):
        i1 = obj.shape[d]
        i2 = obj.maxshape[d]
        s = str(i1)
        if not i2:
            s += "|inf"
        elif i2!=i1:
            s += "|" + str(i2)
        out.append(s)
    return ",".join(out)

def repvalues(obj):
    '''Show a few representative values from the given dataset.'''
    rank = len(obj.shape)
    ret = ""
    if   rank==1:
        ret += str(obj[0])
        if obj.shape[0]>1:
            ret += " " + str(obj[1])
        if obj.shape[0]>3:
            ret += " .."
        if obj.shape[0]>2:
            ret += " " + str(obj[-1])
    elif rank==2:
        if obj.shape[0]<2:
            raise Exception( "Unexpected unused dimension" )
        ret += str(obj[0][0])
        if obj.shape[1]>2:
            ret += " .."
        if obj.shape[1]>1:
            ret += " " + str(obj[0][-1])
        ret += " | "
        if obj.shape[0]>2:
            ret += ".. | "
        ret += str(obj[-1][0])
        if obj.shape[1]>2:
            ret += " .."
        if obj.shape[1]>1:
            ret += " " + str(obj[-1][-1])
    else:
        pass # do not attempt to print high-dimensional table
    return ret

def showobject( obj, name, indent ):
    '''Recursive printout for one group or dataset.'''
    # print name
    out = name
    if "NX_class" in obj.attrs:
        out += ":" + obj.attrs["NX_class"]
    print( indent + out )
    ind = indent + INDENT
    # has a similar group been shown before?
    if isinstance( obj, h5py.Group ):
        if len(obj.keys())>0:
            for og in showngroups:
                if samestruct(obj,og):
                    print( ind+"same struct as "+og.name )
                    return
    # print attributes:
    for k,v in obj.attrs.items():
        if k == "NX_class":
            continue
        print( ind + "@" + k + " = " + str(v) )
    # print subgroups or dataset contents:
    if isinstance( obj, h5py.Group ):
        for k,v in obj.items():
            showobject( v, k, ind )
            showngroups.append( v )
    elif isinstance( obj, h5py.Dataset ):
        if obj.dtype.kind == "S" and len(obj.shape)==1 and obj.shape[0]==1:
            print( ind + str(obj[0]) )
        else:
            print( ind + str(obj.dtype) + "[" + shapestr(obj) + "] " + 
                   repvalues(obj) )
    else:
        raise Exception( "Unexpected object type" )

# main

parser = argparse.ArgumentParser(description='Show structure of a NeXus file')
parser.add_argument('file', metavar='file', help='input NeXus file')
args = parser.parse_args()
fname = args.file

f = h5py.File( fname, 'r' )

showngroups = []
    
showobject( f, fname, "" )