"""extracttree.py
   Example of how to read in a merger tree file in HDF5 format
   and to create a table containing all the data for a particular tree.
"""

# Note: only if we pick up whole trees will the indices be guaranteed
# to be consecutive.  As we only have a single tree in our test data set,
# this example takes halo 7 to be the root halo as that does have
# consecutive indices for its subtree.

# If the indices are consecutive then there is no need to use a special
# Tree class and methods to allow tree searching.  However, we do first 
# do this as a check.

#-------------------------------------------------------------------------------
# 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
import h5py

# The tree module contains a class definition for a combined spatial and temporal tree
import tree

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

# The following data set contains a simple example with both a spatial and merger trees.
infile='data/test.hdf5'
# We are going to open read only
mode='r'

#-------------------------------------------------------------------------------
# Open Merger tree file for reading
fid = h5py.File(infile, mode)

# This shows how easy it is to extract attributes from the file
for (name,value) in fid['/'].attrs.items(): print(name,value)

# Let's extract the data that we want from the HDF5 file.

# Define numpy data type for the table rows
# Maybe there is a clever way to do this, but this example will just do it by hand
treedata_dtype=numpy.dtype([
        ('haloID','int64'),
        ('descendant','int32'),
        ('firstProgenitor','int32'),
        ('firstSubhalo','int32'),
        ('hostHalo','int32'),
        ('neighbour','int32'),
        ('nextSibling','int32'),
        ('snapNum','int32'),
        ('density','f8'),
        ('mass','f8')
        ])

# First let's extract the number of halos
gid = fid['/MergerTree']
nHalo=gid.attrs.get('NHalo')
print('nHalo =',nHalo)

# and now the link that we want.  Without the [:] we get a pointer to the array.
haloID=gid.get('HaloID')[:]
descendant=gid.get('DescendantIndex')[:]
firstProgenitor=gid.get('FirstProgenitorIndex')[:]
firstSubhalo=gid.get('FirstSubhaloIndex')[:]
hostHalo=gid.get('HostHaloIndex')[:]
neighbour=gid.get('NeighbourIndex')[:]
nextSibling=gid.get('NextSiblingIndex')[:]
snapNum=gid.get('Snapshot')[:]
density=gid.get('Density')[:]
mass=gid.get('Mass')[:]

# Next we are going to construct a tree from the information contained in the HDF5 file
# This is only needed for trees that are not depth-first ordered, but it's a nice example
# of how to use iterators in python anyway.  It uses the Tree class in the separate tree.py file.

# Initialise tree nodes...
# Trick: create an extra non-existent node to receive references to index -1 (= last index)
nodes=[tree.Tree(halo) for halo in range(nHalo+1)]
nodes[nHalo]=None

# Now we run through each of the halos extracting those pointers that we need for
# efficient tree traversal
for halo in range(nHalo): 
    nodes[halo].load(nodes[firstSubhalo[halo]],
                     nodes[neighbour[halo]],
                     nodes[firstProgenitor[halo]],
                     nodes[nextSibling[halo]],
                     nodes[hostHalo[halo]],
                     nodes[descendant[halo]])

# Let's pick a halo to act as the root of our tree.
# Ideally, we should use a full tree, but we only have one tree in this example.
# However halo 7 has consecutive indices in its subtree, so we'll use that.
root=7
nodes[root].printheader()
nTreeHalo=0
nodeMin=nodes[root].node
nodeMax=-1
for halo in nodes[root].next():
    nTreeHalo+=1
    nodeMin=min(halo,nodeMin)
    nodeMax=max(halo,nodeMax)
    print(nodes[halo])
print('nTreeHalo =',nTreeHalo)
print('nodeMin, nodeMax =',nodeMin,nodeMax)

# Now let's loop over the tree again, recording the data that we want in our table
# Create a numpty structured array to hold the table
treedata=numpy.empty(nTreeHalo,dtype=treedata_dtype)
ihalo=-1
consecutive=True
firstHalo=nodes[root].node
for halo in nodes[root].next():
    ihalo+=1
    if halo != firstHalo+ihalo: consecutive=False
    treedata[ihalo]=(haloID[halo],descendant[halo],firstProgenitor[halo],firstSubhalo[halo],
                     hostHalo[halo],neighbour[halo],nextSibling[halo],
                     snapNum[halo],density[halo],mass[halo])
print('Consecutive halos = ',consecutive)

# If we have a consecutive list of halos, adjust indices to point to the relevant location
# in our new table.
# Setting the maximum to -1 is not really necessary but helps to make things look tidier.
if consecutive:
    numpy.maximum(treedata['descendant']-nodeMin,-1,treedata['descendant']); 
    numpy.maximum(treedata['firstProgenitor']-nodeMin,-1,treedata['firstProgenitor']); 
    numpy.maximum(treedata['firstSubhalo']-nodeMin,-1,treedata['firstSubhalo']); 
    numpy.maximum(treedata['hostHalo']-nodeMin,-1,treedata['hostHalo']); 
    numpy.maximum(treedata['neighbour']-nodeMin,-1,treedata['neighbour']); 
    numpy.maximum(treedata['nextSibling']-nodeMin,-1,treedata['nextSibling']); 

# We have established that we have a consecutive (depth-first ordered) tree,
# so no need to bother with the Tree class again.  Let's just dump the new
# properties in consecutive order as a check.
for ihalo in range(nTreeHalo):
    print(ihalo,treedata[['firstSubhalo','neighbour','firstProgenitor',
                         'nextSibling','hostHalo','descendant']][ihalo])

# Or rather more simply, but in a different order
print(treedata)

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