import numpy as np

from WEAVEify_routines.step_02_common_routines.WEAVEify_function_calls import \
WEAVEify_function_calls_main

from WEAVEify_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       = "./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       = "Flux_Fields.npy"
npy_filename   = inp_path + inp_file
data_dict      = np.load(npy_filename,allow_pickle = True,encoding = 'latin1').item()
wavelength     = data_dict["wavelength"]
flux_all       = data_dict["flux"]
# ---------------------------------------------



# -----------------------------------
# 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])

# Variation in Magnitude of the source
m_source_arr = np.array([19.6,18.1,19.5,19.8,18.7,20.4,21.1,20.6,19.2,21.9])

# Variation in Magnitude of the sky
m_sky_arr = np.array([21.6,20.1,20.5,21.8,21.7,20.6,21.8,20.8,21.3,21.8])

# 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)


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

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


# Type of noise to be added to the spectra
flag_noise  = "OpR3"

# 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]

# Set this variable to y if need to plot the spectrum
flag_make_quick_plot="n"

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


# -------------------------------------------------------
# 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_WEAVEify.py".

    los_indx     = indx - indx_lb + 1
    spec_indx    = spec_indx_arr[indx]
    flux         = flux_all[:,spec_indx]
    spec_id      = spec_id_arr[indx]
    z_em         = z_em_arr[indx]
    m_source     = m_source_arr[indx]
    m_sky        = m_sky_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_WEAVEify.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,\
    wavelength=wavelength,flux=flux,z_em=z_em,spec_indx=spec_indx,\
    m_source=m_source,m_sky=m_sky,out_path=out_path,\
    flag_noise=flag_noise,flag_output_format=flag_output_format,\
    flag_make_quick_plot=flag_make_quick_plot)

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


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






