"""dumptree.py
   Example of how to write a merger tree in HDF5 format.
   This example needs some data which it reads in from an existing HDF file,
   which makes this whole exercise a bit pointless, but pretend that the data
   was created some other way and now needs to be written out.
   This is not complete, but it gives the general idea.
"""

#-------------------------------------------------------------------------------
# Imports

# The following imports the python3 print function into python2;
# It seesm to be quietly ignored in python3
from __future__ import print_function

import sys
import numpy as np

# The merger tree module contains data structures, methods and global variables
import mt

#-------------------------------------------------------------------------------
# Parameters

infile='data/test.hdf5'
outfile='data/dumptest_Python.hdf5'

# For the purposes of this example, we first have to read in some data.
# Actually, the following creates pointers to the data, to be read below.

mt_in=mt.MergerTree(infile, "r")
attributes = mt_in.attrs('/')

# Read in snap tables
snap_in=mt_in.read_snaptable()
snapprop_in=mt_in.read_snapprop()
snap_attributes=mt_in.attrs("Snapshots")

# Read merger tree.
tree_in=mt_in.read_mergertree()
tree_attributes=mt_in.attrs("MergerTree")

# Now we are going to create a new HDF5 file to write out the data
print('Creating output file')
mt_out=mt.MergerTree(outfile,"w")
print('File =',mt_out.fid)

# Write out the top-level attributes
print('In group',mt_out.fid.name,'attributes are:')
for (name,value) in attributes: 
    print('  ',name,value)
    mt_out.fid.attrs.modify(name,value)

# Create MergerTree directory and set attributes
tree=mt_out.fid.create_group("MergerTree")
print('In group',tree.name,'attributes are:')
for (name,value) in tree_attributes: 
    print('name, value =',name,value)
    tree.attrs.modify(name,value)

# Write out merger tree datasets
print('MergerTree datasets:')
for dset in tree_in:
    # extract name, dtype and data from input dataset
    name=tree_in[dset].name
    dtype=tree_in[dset].dtype
    data=tree_in[dset][()]
    attrs=tree_in[dset].attrs.items()
    # and now recreate in the output file
    dset=tree.create_dataset(name=name,dtype=dtype,data=data)
    print('  ',name)
    print('    ',dset[0:9])
    # and set attributes
    for (name,value) in attrs: 
        print('    ',name,value)
        dset.attrs.modify(name,value)

# Create Snapshot directory and set attributes
snap=mt_out.fid.create_group("Snapshots")
print('In group',snap.name,'attributes are:')
for (name,value) in snap_attributes: 
    print('  ',name,value)
    snap.attrs.modify(name,value)

# Write out Snap table
# Tables require use of a complex dtype.  Fortuntately, that's already set up
# us here in the input file.
# Extract data and dtype from input file
name=snap_in.name
dtype=snap_in.dtype
nsnap=len(snap_in)
data=np.empty(nsnap,dtype=dtype)
for irow in range(nsnap): data[irow]=snap_in[irow]
# and recreate in output file
dset=tree.create_dataset(name=name,dtype=dtype,data=data)
print('Snap table:')
print('  name =',dset.name)
print('  dtype =',dset.dtype)
print('  data =',dset[:])

# Write out SnapProp table
# us here in the input file.
# Extract data and dtype from input file
name=snapprop_in.name
dtype=snapprop_in.dtype
nsnapprop=len(snapprop_in)
data=np.empty(nsnapprop,dtype=dtype)
for irow in range(nsnapprop): data[irow]=snapprop_in[irow]
# Rename the Units property to lower case (for consistency).
# Cannot rename in place: have to make a copy of both dtype and data.
temp=dtype.__repr__()[6:-1]
temp=temp.replace("Units","units")
dtypenew=np.dtype(eval(temp))
print('dtypenew=',dtypenew)
datanew=np.empty(nsnapprop,dtype=dtypenew)
for irow in range(nsnapprop): datanew[irow]=data[irow]
# and recreate in output file
dset=tree.create_dataset(name=name,dtype=dtypenew,data=datanew)
print('SnapProp table:')
print('  name =',dset.name)
print('  dtype =',dset.dtype)
print('  data =',dset[:])

# The output file is not yet a complete examle of the merger tree data
# format, but hopefuly that is enough to illustrate how do do things.

# Close the files
mt_in.close()
mt_out.close()

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