"""Demonstration script to plot the star formation histories of galaxies.

It plots the star formation history, reconstructed from stellar formation
times, for central and satellite galaxies in a narrow stellar mass bin.

This example illustrates using the ReadRegion class to load particles within
the high-resolution region of the simulation at z = 0, the virtual
SubhaloIndex attribute to link particles to their subhaloes, and reading
galaxy data from a specified snapshot from the pre-compiled tables.
"""

# Import required packages
import numpy as np
import hydrangea as hy
import matplotlib.pyplot as plt
from astropy.cosmology import Planck13

# Set script parameters
sim_index = 15                  # Which simulation do we want?
snap_index = 29                 # Which snapshot? 29 --> z = 0
log_mstar_range = [10.0, 10.5]  # Min and max log_mstar to consider
min_sat_log_m200 = 13.0         # Min log_M200 of satellites
plotloc = 'sf_history.png'      # Where to save the output plot?

# Prepare the plot
fig = plt.figure(figsize=(4, 4))

# Set up the simulation object for the run we're working with, to easily
# get the relevant file paths
sim = hy.Simulation(index=sim_index)

# Form the (first) file names of the snapshot and subfind catalogue at the
# snapshot we are working with
snapshot_file = sim.get_snapshot_file(snap_index)
subfind_file = sim.get_subfind_file(snap_index)

# -------------------------------------------------------------------------
# First step: load all star particles within 10r200 from the cluster centre
# (i.e. within the high-resolution region of the simulation)
# -------------------------------------------------------------------------

# Set up a reader for the FOF catalogue, specifying that we are only
# interested in the first object (the central cluster). From this, we can
# access all the FOF properties as attributes of <fof_cl>; they are loaded
# when first required. By default, quantities are returned in
# 'astronomically sensible' units (e.g. Mpc, Gyr, km/s).
fof_cl = hy.SplitFile(subfind_file, 'FOF', read_index=0)

# Set up a reader for star particles within 10r200 from the cluster centre
print("\nSet up stars ReadRegion...")
stars = hy.ReadRegion(snapshot_file, 4, fof_cl.GroupCentreOfPotential,
                      10 * fof_cl.Group_R_Crit200, exact=True)

# Load the stellar masses, central/satellite status, and M200 of all subhaloes
# in the target snapshot. For efficiency, we load these from the pre-compiled
# galaxy tables, and 'translate' these back to subhalo order by indexing these
# with the subhaloes' galaxy IDs. We could also have loaded these data (more
# slowly) directly from the Subfind catalogues, via a SplitFile reader
print("\nLoad subhalo masses and satellite flags...")

# Load galaxy properties
log_mstar = hy.hdf5.read_data(sim.fgt_loc, 'Mstar', read_index=snap_index,
                              index_dim=1)
sat_flag = hy.hdf5.read_data(sim.fgt_loc, 'SatFlag', read_index=snap_index,
                             index_dim=1)
log_m200 = hy.hdf5.read_data(sim.fgt_loc, 'M200', read_index=snap_index,
                             index_dim=1)

# Load the galaxy ID for all subhaloes in the current snapshot, to translate
# galaxy to subhalo order
galaxy_ids = hy.hdf5.read_data(sim.spider_loc,
                               f'Subhalo/Snapshot_{snap_index:03d}/Galaxy')
subhalo_log_mstar = log_mstar[galaxy_ids]
subhalo_sat_flag = sat_flag[galaxy_ids]
subhalo_log_m200 = log_m200[galaxy_ids]

# Find stars in central galaxies in the selected mass range.
# SubhaloIndex is a virtual data set that is created when first accessed
print("\nFinding stars in central galaxies...")
stars.subfind_file = sim.get_subfind_file(snap_index)
ind_cen = np.nonzero((subhalo_log_mstar[stars.SubhaloIndex] >=
                      log_mstar_range[0]) &
                     (subhalo_log_mstar[stars.SubhaloIndex] <
                      log_mstar_range[1]) &
                     (subhalo_sat_flag[stars.SubhaloIndex] == 0))[0]

# Find stars in satellite galaxies in the selected mass range
print("\nFinding stars in satellite galaxies...")
ind_sat = np.nonzero((subhalo_log_mstar[stars.SubhaloIndex] >=
                      log_mstar_range[0]) &
                     (subhalo_log_mstar[stars.SubhaloIndex] <
                      log_mstar_range[1]) &
                     (subhalo_sat_flag[stars.SubhaloIndex] == 1) &
                     (subhalo_log_m200[stars.SubhaloIndex] >
                      min_sat_log_m200))[0]

# Create a histogram of the formation times of stars in central and
# satellite galaxies. We use the initial mass as weights, to retrieve
# the star formation history
formation_time_cen = hy.aexp_to_time(stars.StellarFormationTime[ind_cen])
formation_time_sat = hy.aexp_to_time(stars.StellarFormationTime[ind_sat])
hist_cen, edges = np.histogram(formation_time_cen, bins=28, range=[0, 14],
                               weights=stars.InitialMass[ind_cen],
                               density=True)
hist_sat, edges = np.histogram(formation_time_sat, bins=28, range=[0, 14],
                               weights=stars.InitialMass[ind_sat],
                               density=True)

# Plot the histograms
central_points = edges[:-1] + (edges[1:]-edges[:-1])/2
plt.plot(central_points, hist_cen, color='blue', linestyle='-',
         label='Centrals')
plt.plot(central_points, hist_sat, color='red', linestyle=':',
         label='Satellites')

# Some embellishments on the plot
plt.legend(fontsize=8)

ax = plt.gca()
age_z0 = hy.aexp_to_time(1.0)
ax.set_xlim((0, age_z0))
ax.set_ylim((0, 0.3))
ax.set_xlabel('Age of the Universe [Gyr]')
ax.set_ylabel(r'Star formation history')

# Add redshifts along the upper x-axis
ax2 = ax.twiny()
zred = np.array((4, 2, 1, 0.5, 0.3, 0.1, 0))
age = Planck13.age(zred).value
ax2.set_xlim((0, age_z0))
ax2.set_xticks(age)
ax2.set_xticklabels([str(_z) for _z in zred])
ax2.set_xlabel('Redshift')

# Save the figure
plt.subplots_adjust(left=0.17, bottom=0.15, right=0.96, top=0.88)
plt.show
plt.savefig(plotloc, transparent=False, dpi=150)