import numpy as np

from ASTRO_MOCKS_routines.step_02_common_routines.WEAVEify_function_calls import \
WEAVEify_function_calls_main

from ASTRO_MOCKS_routines.step_02_common_routines.mpi_init import \
distribute_jobs_per_processor


# ---------------------------------------------
# Read inputs ie.e, cosmological density, velocity and temperature fields
# for WEAVE test example
# ---------------------------------------------
inp_path      = "./ASTRO_MOCKS_routines/ASTRO_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/INPUT/"
inp_file      = "Cosmological_Mock_Fields.npy"
npy_filename  = inp_path + inp_file
data_dict     = np.load(npy_filename,allow_pickle = True,encoding = 'latin1').item()
z_arr         = data_dict["z_arr"]
Delta_arr_all = data_dict["Delta_arr"]
T_gas_arr_all = data_dict["T_gas_arr"]
v_gas_arr_all = data_dict["v_gas_arr"]
f_HI_arr_all  = data_dict["f_HI_arr"]
# ---------------------------------------------



# -----------------------------------
# Parameter Variation Block Starts
# Here users can vary these parameters by (i) explicitily specifying,
# or (ii) using some probablity distribution functions
# or (iii) read from files



# Variation in Emission Redshift
z_em_arr = np.array([3.14,2.83,2.34,2.47,2.97,3.07,3.46,3.71,3.24,2.20])

# Index variable that varies sightlines for test cosmological mocks
spec_indx_arr = np.arange(z_em_arr.shape[0],dtype=np.int32)

# The following variable assigns unique name to Mock spectra while saving the
# output. Note that spec_id_arr need not be integers but could also be
# array of strings
spec_id_arr  = np.arange(z_em_arr.shape[0],dtype=np.int32)

# DLA_parameter log_N_HI
log_N_HI_DLA_arr  = np.array([None,18.1,20.8,21.7,19.4,19.2,20.5,21.3,20.1,18.8])

# DLA_parameter Redshift of DLA
z_DLA_arr = np.array([None,2.53,2.20,2.27,2.77,2.90,3.16,3.61,3.04,2.10])

# DLA_parameter Line width parameter for DLA
b_DLA_arr = np.array([None,5.0,8.0,9.0,10.0,7.0,8.0,9.4,7.6,8.2])

# Output path where results will be stored
out_path = "./ASTRO_MOCKS_routines/ASTRO_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/OUTPUT/PARALLEL_MPI_MULTIPLE_SPECTRA/"

# Tell the code to calculate the optical depth for HI forest and DLA
flag_tau = "HI_DLA" #None #"HI_DLA"

# Tell the code to use HI forest and DLA optical depths to calculate flux
flag_absorber = "HI_DLA"

# Tell the code what output format to be use fits or hdf5
flag_output_format = "fits"


# Tells code to use parallelize the job or not
flag_mpi    = "y"

# Total number of spectra to be generated
N_spectra   = spec_id_arr.shape[0]


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


# -------------------------------------------------------
# Users do not need to change following lines of the code
# This part of code initilizes MPI variables if flag_mpi="y"
# and decides how to parallelize the job

mpi_dict    = distribute_jobs_per_processor(flag_mpi,N_spectra)
rank        = mpi_dict["rank"]
comm        = mpi_dict["comm"]
N_proc      = mpi_dict["N_proc"]
indx_lb     = mpi_dict["indx_lb"]
indx_ub     = mpi_dict["indx_ub"]
N_los_total = mpi_dict["N_los_total"]
# -------------------------------------------------------

# Put for loop here to automate script for many WEAVE mocks
# One can use as many nested for loops as necessary

for indx in range(indx_lb,indx_ub):

    # -------------------------------------------------------------------
    # Please update all the variables that you want to vary from sightline
    # to sightline. Users can vary as many parameters as listed in file
    # "default_parameters_astro_mocks.py".

    los_indx     = indx - indx_lb + 1
    spec_indx    = spec_indx_arr[indx]
    Delta_arr    = Delta_arr_all[:,spec_indx]
    T_gas_arr    = T_gas_arr_all[:,spec_indx]
    v_gas_arr    = v_gas_arr_all[:,spec_indx]
    f_HI_arr     = f_HI_arr_all[:,spec_indx]
    spec_id      = spec_id_arr[indx]
    z_em         = z_em_arr[indx]
    log_N_HI_DLA = log_N_HI_DLA_arr[indx]
    z_DLA        = z_DLA_arr[indx]
    b_DLA        = b_DLA_arr[indx]

    # Users can add as many variables as they need to vary for mock spectra

    # VERY IMPORTANT:
    # Note that sequence of keyword arguments in following function is not
    # important. BUT make sure that the variable names are exactly the same as
    # the variable (parameter) names listed in "default_parameters_astro_mocks.py"
    # Otherwise the default value of variables will be used by the code.

    # ALWAYS ADD flag_mpi, rank, N_proc, N_los_total, los_indx, spec_id 
    # variables in argument list. They are necessary for MPI routines.
    # Rest of the variables are optional

    WEAVEify_function_calls_main(flag_mpi=flag_mpi,rank=rank,N_proc=N_proc,\
    N_los_total=N_los_total,los_indx=los_indx,spec_id=spec_id,\
    z_arr=z_arr,Delta_arr=Delta_arr,T_gas_arr=T_gas_arr,v_gas_arr=v_gas_arr,\
    f_HI_arr=f_HI_arr,z_em=z_em,spec_indx=spec_indx,\
    log_N_HI_DLA=log_N_HI_DLA,z_DLA=z_DLA,b_DLA=b_DLA,\
    flag_tau=flag_tau,flag_absorber=flag_absorber,out_path=out_path,\
    flag_noise=flag_noise,flag_output_format=flag_output_format)

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


# Synchronize the code at  this point
if flag_mpi.upper()=="Y":comm.Barrier()






