import numpy as np

from scipy.interpolate import RegularGridInterpolator

def cloudy_table_interpolation_routines_main(inp_dict):
    cloudy_table_dict  = read_cloudy_table(inp_dict)
    cloudy_table_dict  = generate_intrpolation_function(cloudy_table_dict)
    cloudy_intrp_funct = cloudy_table_dict["cloudy_intrp_funct"]
    return cloudy_intrp_funct


def read_cloudy_table(inp_dict_org):
    inp_dict           = inp_dict_org.copy()
    inp_path_aux_files = inp_dict["inp_path_aux_files"]
    inp_path_cloudy    = inp_path_aux_files + "/METAL_ADDITION/METHOD_1_CLOUDY_INTERPOLATION/"
    inp_file_cloudy    = "Cloudy_Ionization_Fraction_Tables.npy"

    zmean_val             = inp_dict["zmean_val"]
    metal_specie_ion      = inp_dict["metal_specie_ion"]
    metal_specie_ion_list = [metal_specie_ion]
    npy_filename          = inp_path_cloudy + inp_file_cloudy

    #CHECK_AND_INITIALIZE_KEY_VALUE(inp_dict,"specie_ion_list",specie_ion_list)
    block_dict        = generate_cloudy_table_block_dict(inp_dict,metal_specie_ion_list)
    log_Z_block       = block_dict["log_Z_block"]
    log_nH_block      = block_dict["log_nH_block"]
    log_T_block       = block_dict["log_T_block"]
    log_f_XI_block    = block_dict["log_f_" + metal_specie_ion + "_block"]

    cloudy_intrp_dict = np.load(npy_filename,allow_pickle=True,encoding='latin1').item()
    log_Z_arr         = cloudy_intrp_dict[log_Z_block]
    log_nH_arr        = cloudy_intrp_dict[log_nH_block]
    log_T_arr         = cloudy_intrp_dict[log_T_block]
    log_f_XI_arr      = cloudy_intrp_dict[log_f_XI_block]
    #print log_f_XI_arr[0,0,1]
    #print log_nH_arr[0],log_nH_arr[-1]
    #print log_T_arr[0],log_T_arr[-1]
    #print log_Z_arr[0],log_Z_arr[1]
    cloudy_table_dict = {"log_Z_arr":log_Z_arr,"log_nH_arr":log_nH_arr,\
                         "log_T_arr":log_T_arr,"log_f_XI_arr":log_f_XI_arr}
    return cloudy_table_dict

def generate_intrpolation_function(cloudy_table_dict):
    log_nH_arr         = cloudy_table_dict["log_nH_arr"]
    log_T_arr          = cloudy_table_dict["log_T_arr"]
    log_Z_arr          = cloudy_table_dict["log_Z_arr"]
    log_f_XI_arr       = cloudy_table_dict["log_f_XI_arr"]
    cloudy_intrp_funct = RegularGridInterpolator((log_nH_arr, log_T_arr, log_Z_arr), log_f_XI_arr)
    cloudy_table_dict.update({"cloudy_intrp_funct":cloudy_intrp_funct})
    return cloudy_table_dict

def generate_cloudy_table_block_dict(inp_dict,metal_specie_ion_list):
    uvb_model_metal       = inp_dict["uvb_model_metal"]
    zmean_val             = round(inp_dict["zmean_val"],4)
    common_data_block     = "/" + uvb_model_metal + "/Common_Data_Grids/"
    log_nH_block          = common_data_block + "log_nH_Grids"
    log_Z_block           = common_data_block + "log_Z_Grids"
    log_T_block           = common_data_block + "log_T_Grids"
    block_dict            = {"log_nH_block":log_nH_block,"log_T_block":log_T_block,\
                             "log_Z_block":log_Z_block}
    for metal_specie_ion in metal_specie_ion_list:
        key_name = "log_f_" + metal_specie_ion + "_block"
        key_val  = "/" + uvb_model_metal + "/Specie_Fraction/z-" + str("%.4f"%zmean_val) + "/log_f_" + metal_specie_ion 
        block_dict.update({key_name:key_val})
    return block_dict
