import numpy as np
import matplotlib.pyplot as plt

from scipy.interpolate import interp1d

from WEAVEify_routines.step_02_common_routines.dictionary_functions import READ_DICTIONARY_FROM_FILE

from WEAVEify_routines.step_02_common_routines.main_routines_tau import extract_specie_list

from WEAVEify_routines.step_02_common_routines.wavelength_array_calculation import \
wavelength_array_calculation_main

def generate_flux_main(inp_dict):
    inp_dict          = wavelength_array_calculation_main(inp_dict)
    input_type_str    = str(inp_dict["input_type"])
    if input_type_str == "1":
        z_em              = inp_dict["z_em"]
        line_dict         = inp_dict["line_dict"]
        lambda_rest_HI    = line_dict["HI-1215"]["lambda_rest"]
        z_sim_arr_all     = inp_dict["z_arr"]
        bool_arr_z_em     = z_sim_arr_all <= z_em
        z_sim_arr         = z_sim_arr_all[bool_arr_z_em]
        lambda_arr_z      = lambda_rest_HI * (1.0 + z_sim_arr)
        N_z               = z_sim_arr.shape[0]

        wavelength_arr    = inp_dict["lambda_arr_obs"]

        tau_total_arr       = np.zeros(wavelength_arr.shape[0])
        tau_total_HI_arr    = np.zeros(wavelength_arr.shape[0])
        tau_total_metal_arr = np.zeros(wavelength_arr.shape[0])
        tau_total_DLA_arr   = np.zeros(wavelength_arr.shape[0])

        if "HI" in inp_dict["flag_absorber"].upper():
            tau_total_HI_arr    = collect_optical_depth_for_given_specie_type(inp_dict,"HI")

        if "METAL" in inp_dict["flag_absorber"].upper():
            tau_total_metal_arr = collect_optical_depth_for_given_specie_type(inp_dict,"metal")

        if "DLA" in inp_dict["flag_absorber"].upper():
            tau_total_DLA_arr = collect_optical_depth_for_given_specie_type(inp_dict,"DLA")

        tau_total_arr       = tau_total_HI_arr + tau_total_metal_arr + tau_total_DLA_arr

    elif input_type_str == "2":
        tau_total_arr  = inp_dict["tau_arr"]
        wavelength_arr = inp_dict["lambda_arr_obs"]

    Flux_total_arr      = np.exp(-1.0 * tau_total_arr)
    #print(input_type_str)
    #plt.plot(wavelength_arr,Flux_total_arr,lw=0.5)
    #plt.plot(wavelength_arr,np.exp(-tau_total_DLA_arr))
    #plt.ylim(-0.1,1.2)
    #plt.show()
    #exit()

    inp_dict_Flux       = {"tau_total_arr":tau_total_arr,\
                           "wavelength_total_arr":wavelength_arr,\
                           "Flux_total_arr":Flux_total_arr}

    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})

    return inp_dict


def collect_optical_depth_for_given_specie_type(inp_dict,specie):


    if "inp_dict_tau" in inp_dict.keys():
        inp_dict_tau         = inp_dict["inp_dict_tau"]
        flag_tau             = inp_dict["flag_tau"]
        if specie.upper() not in flag_tau.upper():
            tau_total_specie_arr = read_data_from_files(inp_dict,specie)
        else:
            tau_total_specie_arr = inp_dict_tau["tau_total_" + specie + "_arr"]

    else:
        tau_total_specie_arr = read_data_from_files(inp_dict,specie)


    return tau_total_specie_arr

def read_data_from_files(inp_dict,specie):
    from WEAVEify_routines.step_11_read_write_mock_spectra.\
    input_output_all_format  import read_optical_depths_all_format

    inp_dict_tau         = read_optical_depths_all_format(inp_dict,flag_absorber = specie)
    tau_total_specie_arr = inp_dict_tau["tau_total_arr"]

    return tau_total_specie_arr

