import numpy as np
import sys

sys.dont_write_bytecode = True


from ASTRO_MOCKS_routines.step_02_common_routines.astro_mocks_function_calls import \
astro_mocks_function_calls_main

# ---------------------------------------------
# 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/"
out_path      = "./ASTRO_MOCKS_routines/ASTRO_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/OUTPUT/SERIAL_SINGLE_SPECTRUM/"
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"]

# Changes for metal profile addition


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

N_los_total   = Delta_arr_all.shape[1]
los_indx      = 0
Delta_arr     = Delta_arr_all[:,los_indx]
T_gas_arr     = T_gas_arr_all[:,los_indx]
v_gas_arr     = v_gas_arr_all[:,los_indx]
f_HI_arr      =  f_HI_arr_all[:,los_indx]

# Note that sequence of keyword arguments 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.

astro_mocks_function_calls_main(\
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,\
out_path=out_path,\
spec_id=los_indx
)





