import numpy as np

from WEAVEify_routines.step_02_common_routines.WEAVEify_function_calls_metal_profile import \
WEAVEify_function_calls_main

# ---------------------------------------------
# Read inputs ie.e, cosmological density, velocity and temperature fields
# for WEAVE test example
# ---------------------------------------------
inp_path      = "./WEAVEify_routines/WEAVE_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/INPUT/"
out_path      = "./WEAVEify_routines/WEAVE_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
flag_metal_method = "metal_profile"
flag_absorber     = "HI_metal"


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

spec_indx     = 0
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]

cont_rand_seed   = 123
#flag_noise       = None
#flag_cont_method = PCA

# Make plots for test examples
flag_make_quick_plot = "y"

# 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_WEAVEify.py"
# Otherwise the default value of variables will be used by the code.

WEAVEify_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,\
cont_rand_seed=cont_rand_seed,out_path=out_path,\
flag_make_quick_plot=flag_make_quick_plot,\
flag_metal_method=flag_metal_method,\
flag_absorber=flag_absorber)





