"""
Short example script to demonstrate how to read in FOF and Subhalo properties.
"""

import sim_tools as st       # Set of helper functions for reading in data
import numpy as np
from pdb import set_trace

# Specify directory of the simulation. Here, we use the Hydro version of CE-0
simdir = '/net/quasar/data3/Hydrangea/CE-0/HYDRO/'

# Specify one of the HDF5 files that you want to read in:
subdir = simdir + '/data/groups_029_z000p000/eagle_subfind_tab_029_z000p000.0.hdf5'

# ... or, more easily, use a shortcut function:
# (returns exactly the same file path, but less effort!)
subdir = st.form_files(simdir, isnap = 29, types = 'sub')

# N.B.: specify the snapshot number as argument to 'isnap', and the type of file to load as 'types':
#    'sub'     --> eagle_subfind_tab
#    'snap'    --> snapshot
#    'subpart' --> eagle_subfind_particles (like snapshot, but only for particles in FOF groups)
#                  (N.B.: not currently on Leiden system for most snapshots)

# Now load the (total) mass of all subhaloes:
# The `eagleread' function internally combines all the individual output files together.
# The first argument is the first of the files to read from.
# The second argument specifies the data set to read.

mass_plus = st.eagleread(subdir, 'Subhalo/Mass', astro = True)

# N.B.: 'astro = True' makes the script automatically apply the appropriate power of h and a (expansion factor)
#      --> gives quantity in 'astronomical' units (mass: 10^10 M_sun, length: Mpc, velocity: km/s)

mass = mass_plus[0]
# with astro = True (default), eagleread returns a list of three elements: the actual data [0], the 
#       conversion factor applied to the raw data, and the expansion factor of the universe at the time of the 
#       snapshot.  

conv_astro = mass_plus[1]
aexp = mass_plus[2]

# conv_astro should be equal to 1/0.6777, aexp to 1 (z = 0)

# Load masses (M200crit) of all FOF groups:
m200c = st.eagleread(subdir, 'FOF/Group_M_Crit200')[0]

# Find all subhaloes that are part of the first (=most massive) FOF group:
# (use astro=False because these data have no conversion factors)
firstSubhalo = st.eagleread(subdir, 'FOF/FirstSubhaloID', astro = False)
numOfSubhaloes = st.eagleread(subdir, 'FOF/NumOfSubhaloes', astro = False)

# The subhaloes are arranged in order of their FOF group in the Subhalo/... data sets.
# For FOF i, the corresponding subhaloes are firstSubhalo[i] up to and including
#   firstSubhalo[i]+numOfSubhaloes[i]-1
ind_subhalo = np.arange(firstSubhalo[0], firstSubhalo[0]+numOfSubhaloes[0])


set_trace()  # Starts the debugger, so you can have a look at the individual variables