<reponame>muntazirabidi/boss-sbi
'''

script for constructing HOD catalogs for the Quijote LHC simulations
with a forward model of the BOSS LOWZ SGC sample


'''
import os, sys 
import numpy as np 
from simbig import halos as Halos
from simbig import galaxies as Galaxies
from simbig import forwardmodel as FM

np.random.seed(1)

m_hod = sys.argv[1] # hod model 
i0 = int(sys.argv[2]) 
i1 = int(sys.argv[3])
zred = float(sys.argv[4])


def sample_HOD(m_hod): 
    ''' sample HOD value based on priors set by Parejko+(2013)
    '''
    if m_hod == 'zheng07':
        hod_min = np.array([13.2, 0.4, 13.1, 14., 0.7])
        dhod = np.array([0.15, 0.1, 0.4, 0.3, 0.4])
        _hod = hod_min + dhod * np.random.uniform(size=(5))
        return {'logMmin': _hod[0], 'sigma_logM': _hod[1], 'logM0': _hod[2], 'logM1': _hod[3], 'alpha': _hod[4]}
    elif m_hod == 'zheng07_ab': 
        hod_min = np.array([13.2, 0.4, 13.1, 14., 0.7, 0., 0.])
        dhod = np.array([0.2, 0.1, 0.5, 0.4, 0.4, 0.5, 0.5])
        _hod = hod_min + dhod * np.random.uniform(size=(7))
        return {'logMmin': _hod[0], 'sigma_logM': _hod[1], 'logM0': _hod[2], 'logM1': _hod[3], 'alpha': _hod[4], 
                'mean_occupation_centrals_assembias_param1': _hod[5], 
                'mean_occupation_satellites_assembias_param1': _hod[6]}
    else: 
        raise NotImplementedError 


dat_dir = '/tigress/chhahn/simbig/chirag/'
for i_lhc in range(i0, i1):
    print('LHC %i' % i_lhc)
    # read in halo catalog
    halos = Halos.Quijote_LHC_HR(i_lhc, z=zred)
    for i_hod in range(5): 
        print('  HOD %i' % i_hod)
        fgal = os.path.join(dat_dir, 'hod.quijote_LH%i.z%s.%s.%i.bf' % 
                (i_lhc, str(zred).replace('.', 'p'), m_hod, i_hod))
        fhod = fgal.replace('.bf', '.npy')

        if not os.path.isfile(fhod) or not os.path.isdir(fgal): 
            # sample HOD parameters
            theta_hod = sample_HOD(m_hod)
            np.save(fhod, theta_hod)
            # populate box with HOD
            hod = Galaxies.hodGalaxies(halos, theta_hod, seed=0, hod_model=m_hod)
            # save box 
            hod.save(fgal) 
