#! /usr/bin/python

# Create an HDF5 format merger tree from a sussing format using the h5py interface

import h5py
#import sys
import numpy
#import glob
#import re

import myglobals
import smt2hdf

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

# Parameters

# Debug/verbosity/testing
myglobals.debug=False
myglobals.verbosity=2
myglobals.test=False

# Tree style
treeStyle='Rockstar'
#treeStyle='SUBFIND_ALL'

# Simulation
#simulation='62.5'
simulation='125'

# Input:
if myglobals.test: title="Test Merger Tree"
else: title="nIFTy merger tree"
if myglobals.test: description="A sample test merger tree created from mock data in SMT format"
else: description="The nIFTy "+treeStyle+" ascii trees converted to HDF5 format"
version=0
subversion=1
if simulation=='62.5':
    h100=0.704
    boxsizeMpc=62.5
    omegaBaryon=0.045
    omegaCDM=0.272-omegaBaryon
    omegaLambda=0.728
    sigma8=0.807
elif simulation=='125':
    h100=0.6777
    boxsizeMpc=125.
    omegaBaryon=0.048252
    omegaCDM=0.307115-omegaBaryon
    omegaLambda=0.692885
    sigma8=0.8288
    
# Location of table of snapshot info, halo files, merger tree, output file
if myglobals.test: 
    snaplistfile="data/data_snaplist.txt"
    snaphalofiles="data/*test.AHF_halos"
    mergertreefile="data/test.tree"
    outfile = "data/test.hdf5"
elif simulation=='62.5':
    snaplistfile="nifty_data_62.5/data_snaplist.txt"
    snaphalofiles="nifty_data_62.5/*.AHF_halos"
    mergertreefile="nifty_data_62.5/MergerTree"+treeStyle+".txt"
    outfile="nifty_data_62.5/MergerTree"+treeStyle+".hdf5"
elif simulation=='125':
    snaplistfile="nifty_data_125/sussing_index.lst"
    snaphalofiles="nifty_data_125/*.AHF_halos"
    mergertreefile="nifty_data_125/sussing_tree.list"
    outfile="nifty_data_125/MergerTree.hdf5"

# Debug using a subset of the halos
if myglobals.debug:
    snaphalofiles="nifty_debug/*.AHF_halos"
    mergertreefile="nifty_debug/MergerTree.txt"
    outfile="nifty_debug/MergerTree.hdf5"

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

# Open the file for output
f = h5py.File(outfile, "w")

# Add root ("/") attributes
f.attrs["Title"] = title
f.attrs["Description"] = description
f.attrs["Version"] = version
f.attrs["Subversion"] = subversion
f.attrs["H100"] = h100
f.attrs["BoxsizeMpc"] = boxsizeMpc
f.attrs["OmegaBaryon"] = omegaBaryon
f.attrs["OmegaCDM"] = omegaCDM
f.attrs["OmegaLambda"] = omegaLambda
f.attrs["Sigma8"] = sigma8

# Read in snapshot info; create and write Snapshots group and tables
snapgroup = f.create_group("Snapshots")
if myglobals.verbosity > 0: print('Reading snaptable...')
smt2hdf.loadsnaptable(snapgroup, snaplistfile)
if myglobals.verbosity > 0: print('...done.')

# Read in halos associated with each snapshot; create appropriate subgroups and write out
# Return information that we need later to populate MergerTree group
if myglobals.verbosity > 0: print('Reading snapfiles...')
hostHalo_dict,mass_dict = smt2hdf.loadsnapfiles(snapgroup, snaphalofiles)
if myglobals.verbosity > 0: print('...done.')

# Load the Sussing Merger Tree format Merger tree data
if myglobals.verbosity > 0: print('Reading mergertree...')
progenitor_dict = smt2hdf.loadmergertree(mergertreefile)
if myglobals.verbosity > 0: print('...done.')

# create group /MergerTree to hold output and set some attributes
mergerTree = f.create_group("MergerTree")
mergerTree.attrs["NHalo"] = len(progenitor_dict)
mergerTree.attrs["TableFlag"] = 0                 # Store properties as separate arrays
mergerTree.attrs["HaloIndexOffset"] = 0           # Pointers have C/Python-style numbering

# Create ordered list of haloIDs
if myglobals.verbosity > 0: print('Sorting halos...')
haloids = numpy.array(list(progenitor_dict.keys()), dtype=int)
haloids.sort()
if myglobals.verbosity > 0: print('...done.')

# Create tree pointers and write tree to MergerTree group
if myglobals.verbosity > 0: print('Populating merger trees...')
smt2hdf.populatemergertree(mergerTree, haloids, progenitor_dict, mass_dict, hostHalo_dict)
if myglobals.verbosity > 0: print('...done.')

# Close HDF output file
f.close()
