"""mt.py
   Merger tree class to simplfy some i/o functions for merger trees.
   It's not entirely obvious that a separate class is needed 
   - could just put these calls directly into the main program.
"""

#-------------------------------------------------------------------------------
import h5py

#-------------------------------------------------------------------------------
class MergerTree:
    
    def __init__(self, file=None, mode='r'):
        """Initialise the class, opening the file with the mode given"""
        if file: 
            self.open(file, mode)
        else:
            print('mt usage: mt(file, mode=mode)')

    # Return set of attributes associated with object
    def attrs(self, object):
        return self.fid[object].attrs.items()

    # Close the HDF5 file
    def close(self):
        self.fid.close()

    # Open the HDF5 file
    def open(self, file, mode):
        self.fid = h5py.File(file, mode)

    # Return a handle to the MergerTree group
    def read_mergertree(self):
        return self.fid["MergerTree"]

    # Return a handle to the Snapshot table
    def read_snaptable(self):
        return self.fid["Snapshots/Snap"]

    # Return a handle to the SnapProp table
    def read_snapprop(self):
        return self.fid["Snapshots/SnapProp"]

#-------------------------------------------------------------------------------
