"""tracetree.py
   Example of how to read in a merger tree in HDF5 format and to use a
   specially-defined Tree class to allow spatial and/or temporal searching
   of the tree.
   Usage: python tracetree.py
"""

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

# The following imports the python3 print function into python2;
# It seems 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)

# Next we are going to construct a tree from the information contained in the HDF5 file

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

# and now the links that we want.  Without the [:] we get a pointer to the array.
firstSubhalo=gid.get('FirstSubhaloIndex')[:]
neighbour=gid.get('NeighbourIndex')[:]
firstProgenitor=gid.get('FirstProgenitorIndex')[:]
nextSibling=gid.get('NextSiblingIndex')[:]
hostHalo=gid.get('HostHaloIndex')[:]
descendant=gid.get('DescendantIndex')[:]
    
# 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]])

# The end main halos are those that have no Descendants and no HostHalo
# In this simple example there is only one such halo
endMainHalos=numpy.where((descendant==-1) & (hostHalo==-1))[0]
print(endMainHalos)

# Let's trace all the halos in the combined spatial and temporal tree
nodes[0].printheader()
for halo in nodes[0].next():
    print(nodes[halo])
    
# Here's a trace of just the spatial tree
nodes[0].printheader()
for halo in nodes[0].next('spatial'):
    print(nodes[halo])

# and here are the temporal trees for each end halo
endHalos=numpy.where(descendant==-1)[0]
for endHalo in endHalos:
    nodes[endHalo].printheader()
    for halo in nodes[endHalo].next('temporal'):
        print(nodes[halo])

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