from __future__ import print_function
import numpy as np
import sys,os

from ASTRO_MOCKS_routines.step_02_common_routines.dictionary_functions import  \
CHECK_AND_INITIALIZE_KEY_VALUE,CHECK_IF_KEY_PRESENT

from ASTRO_MOCKS_routines.step_02_common_routines.directory_functions import \
MAKE_DIRECTORY

def astro_mocks_function_calls_main(**kwargs):
    inp_dict_script = locals()["kwargs"]

    from default_parameters_astro_mocks import step_a_initialize_astro_mocks_parameter,\
    step_b_user_input_command_line,step_d_initialize_dependent_variables,\
    step_c_user_input_script

    from ASTRO_MOCKS_routines.step_03_optical_depth_calculation_HI_forest.\
    step_03_HI_forest import step_03_HI_forest_main

    from ASTRO_MOCKS_routines.step_05_optical_depth_calculation_DLA.step_05_DLA_addition import\
    step_05_DLA_addition_main

    from ASTRO_MOCKS_routines.step_04_optical_depth_calculation_METALS.step_04_metal_addition import \
    step_04_metal_addition_main

    from ASTRO_MOCKS_routines.step_07_read_write_mock_spectra.input_output_all_format \
    import save_optical_depths_all_format,save_flux_fields_all_format

    from ASTRO_MOCKS_routines.step_06_generate_flux_for_all_species.generate_flux import \
    generate_flux_main

    from ASTRO_MOCKS_routines.step_02_common_routines.echo_functions import \
    process_completion_print_function

    from ASTRO_MOCKS_routines.step_01_astro_mock_utils.util_03_plot_spectra.make_quick_plots \
    import make_quick_plots_main


    inp_dict        = step_a_initialize_astro_mocks_parameter()
    inp_dict        = step_c_user_input_script(inp_dict,inp_dict_script)
    inp_dict        = step_b_user_input_command_line(inp_dict)
    inp_dict        = step_d_initialize_dependent_variables(inp_dict)

    process_completion_print_function(inp_dict,"Starting mock making code")

    flag_mpi = inp_dict["flag_mpi"]
    out_path = inp_dict["out_path"]

    if inp_dict["rank"] == 0: MAKE_DIRECTORY(out_path) # Make output directory where data would be saved


    flag_tau       = inp_dict["flag_tau"]
    flag_absorber  = inp_dict["flag_absorber"]

    if (flag_tau is None):
        flag_tau = "NONE"

    if (flag_absorber is None):
        flag_absorber = "NONE"


    if "HI" in flag_tau.upper():
        inp_dict = step_03_HI_forest_main(inp_dict)

    if "METAL" in flag_tau.upper():
        inp_dict = step_04_metal_addition_main(inp_dict,flag_flux_method="n")

    if "DLA" in flag_tau.upper():
        inp_dict = step_05_DLA_addition_main(inp_dict)

    if flag_tau != "NONE":
        inp_dict = save_optical_depths_all_format(inp_dict)

    if flag_absorber != "NONE":
        inp_dict = generate_flux_main(inp_dict)
        inp_dict = save_flux_fields_all_format(inp_dict)

    flag_make_quick_plot = inp_dict["flag_make_quick_plot"]

    if flag_make_quick_plot.upper() == "Y":
        inp_dict  = make_quick_plots_main(inp_dict)
        process_completion_print_function(inp_dict,"Quick plot creation completed")
        process_completion_print_function(inp_dict,"Plot Location : " + inp_dict["out_path"] + inp_dict["out_file_plot"])

    process_completion_print_function(inp_dict,"Astro-Mock making code all processes completed")


    return inp_dict
