"""
Example code to generate a stellar mass (surface density) image
with SPH smoothing (optional).

Written 16-Nov-2018
"""

import sys
sys.path.insert(0, '../')

import numpy as np
import sim_tools as st

from astropy.io import ascii

import yb_utils as yb
import hydrangea_tools as ht
import image_routines as ir

import matplotlib
matplotlib.use('pdf')
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
matplotlib.rcParams['font.family'] = 'serif' 
matplotlib.rcParams['font.serif'][0] = 'palatino'

from pdb import set_trace

# ================= System-specific settings ====================================

# Specify the base directory of the simulation to process
rundir = '/virgo/simulations/Hydrangea/10r200/CE-21/HYDRO/'

# Specify the output image file (without extension):
plotloc = '/virgo/scratch/ybahe/HYDRANGEA/IMAGES/TESTS/StellarImage_CE21_BCG_snap'


# ================= Image settings ==============================================

mark_galaxies = False    # Optionally overplot location of galaxies
mark_dead_galaxies = False
save_maps = False     # Store images as HDF5 as well as picture
desNGB = 32          # Number of neighbours for smoothing calculation
plot_snap = 29       # Which snapshot to plot

ptype = 4            # Particle type to plot (stars = 4)
imsize = 1.0         # (Half-)image size in pMpc
min_mpeak_otherGal = 7.0
fixedSmoothingLength = 0  # Set to > 0 to disable variable smoothing length computation
numPix = 1000        # Image side length in pixels
camDir = [0, 0, -1]  # Specify viewing direction, default = down z-axis

# ===============================================================================
# ===============================================================================


posloc = rundir + '/highlev/GalaxyPositionsSnap.hdf5'
spiderloc = rundir + '/highlev/SpiderwebTables.hdf5'
fgtloc = rundir + '/highlev/FullGalaxyTables.hdf5' 

snaplistloc = rundir + "/sneplists/allsnaps.dat"
snaplist = ascii.read(snaplistloc)
aexpSnap = np.array(snaplist['aexp'])

galID = yb.read_hdf5(spiderloc, 'Subhalo/Snapshot_029/Galaxy')[0]  # Center on central cluster

lastsnap = yb.read_hdf5(spiderloc, 'LastSnap')
firstsnap = yb.read_hdf5(spiderloc, 'FirstSnap')
mpeak = yb.read_hdf5(fgtloc, 'Mstar')[:, plot_snap]

snapStr = 'Snapshot_' + str(plot_snap).zfill(3)

pos_allGal = yb.read_hdf5(posloc, "Centre")[:, plot_snap, :]
pos_gal = pos_allGal[galID, :]

# Slightly lazy way to set comoving <-> proper conversion factor
aexp_factor = aexpSnap[plot_snap]
h_factor = 1/0.6777
conv_astro_pos = h_factor*aexp_factor

print("Determined galaxy position as {:.3f}/{:.3f}/{:.3f}..."
      .format(*pos_gal))

snapdir = st.form_files(rundir, isnap = plot_snap, types = 'snap')

# Set up a mask to efficiently read only the required region from the simulation
readReg = ht.ReadRegion(snapdir, ptype, [*pos_gal/conv_astro_pos, imsize*np.sqrt(3)/conv_astro_pos])

# Read position and mass of particles in target region
pos = readReg.read_data("Coordinates", astro = True)

if ptype != 1:
    mass = readReg.read_data("Mass", astro = True)
else:
    mass = np.zeros(pos.shape[0])+st.m_dm(rundir, astro = True)

deltaPos = pos - pos_gal[None, :]

# Select only those particles that are actually within target sphere
rad = np.linalg.norm(deltaPos, axis = 1)
ind_sphere = np.nonzero(rad < imsize*np.sqrt(3))[0]
mass = mass[ind_sphere]
pos = pos[ind_sphere, :]

# Find galaxies in the field-of-view
ind_in_field = np.nonzero((np.max(np.abs(pos_allGal-pos_gal[None, :]), axis = 1) < imsize) & (mpeak >= min_mpeak_otherGal) & (aexpSnap[lastsnap] >= aexp_factor) & (aexpSnap[firstsnap] <= aexp_factor))[0]

if mark_dead_galaxies:
    ind_in_field = np.nonzero((np.max(np.abs(pos_allGal-pos_gal[None, :]), axis = 1) < imsize) & (mpeak >= min_mpeak_otherGal) & (aexpSnap[firstsnap] <= aexp_factor))[0]
    subind_dead = np.nonzero(aexpSnap[lastsnap[ind_in_field]] > aexp_factor)[0]

pos_in_field = (pos_allGal[ind_in_field, :]-pos_gal[None, :])

if fixedSmoothingLength > 0:
    hsml = np.zeros(len(mass), dtype = np.float32) + fixedSmoothingLength
else:
    hsml = None
        
# Generate actual image
image_weight_all, image_quant = ir.make_sph_image_new_3d(
    pos, mass, mass, hsml, DesNgb=desNGB,imsize=numPix, zpix = 1, boxsize = imsize, 
    CamPos = pos_gal, CamDir = camDir, CamAngle = [0,0,0], CamFOV = [0.0, 0.0], 
    make_deepcopy = True, zrange = [-imsize, imsize], tau = 1e6, return_hsml = False,
    verbose = True)
    
# Display image to check things went fine:
vmin, vmax = 1, 5.5
cmap = plt.cm.Greys_r

plt.figure(figsize = (5.0, 5.0), dpi = 400)
plt.imshow(np.log10(image_weight_all[:, :, 1]+1e-5), cmap = cmap, origin = 'lower', 
           extent = [-imsize, imsize, -imsize, imsize], vmin = vmin, vmax = vmax, 
           interpolation = 'none')

# Some embellishments on the image
plt.text(-0.045/0.05*imsize, 0.045/0.05*imsize, 'z = {:.3f}' 
         .format(1/aexp_factor - 1), va = 'center', 
         ha = 'left', color = 'white')

plt.text(0.045/0.05*imsize, 0.045/0.05*imsize, 'galID = {:d}' 
         .format(galID), va = 'center', 
         ha = 'right', color = 'white')

if mark_galaxies:

    for iigal, igal in enumerate(ind_in_field):

        if aexpSnap[lastsnap[igal]] < aexp_factor:
            color = 'red'
        else:
            color = 'limegreen'

        plt.scatter(pos_in_field[iigal, 0], pos_in_field[iigal, 1], 30, 
                    edgecolor = color, facecolor = 'none', alpha = 0.5)

        plt.text(pos_in_field[iigal, 0]+imsize/100, pos_in_field[iigal, 1]+imsize/100, 
                 str(igal), color = color, va = 'bottom', ha = 'left', alpha = 0.5, fontsize = 6)
        
            
ax = plt.gca()
ax.set_xlabel(r'$\Delta x$ [pMpc]')
ax.set_ylabel(r'$\Delta y$ [pMpc]')
    
ax.set_xlim((-imsize, imsize))
ax.set_ylim((-imsize, imsize))

plt.subplots_adjust(left = 0.15, right = 0.95, bottom = 0.15, top = 0.92)
plt.savefig(plotloc + str(plot_snap).zfill(4) + '.png', dpi=400)
plt.close()

if save_maps:
    maploc = plotloc + str(plot_snap).zfill(4) + '.hdf5'
    yb.write_hdf5(image_weight_all, maploc, 'Image', new = True)
    yb.write_hdf5(np.array((-imsize, imsize, -imsize, imsize)), maploc, 'Extent')
            
    
print("Done!")