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,CHECK_AND_INITIALIZE_KEY_VALUE

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 tau_summation_routines_main(inp_dict,flag_absorber="HI"):
    CHECK_AND_INITIALIZE_KEY_VALUE(inp_dict,"inp_dict_tau",{})

    inp_dict          = wavelength_array_calculation_main(inp_dict)
    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"]
    inp_dict_tau      = inp_dict["inp_dict_tau"]


    if flag_absorber.upper() == "HI":
        tau_total_HI_arr = compute_optical_depth_for_given_specie_type(inp_dict,bool_arr_z_em,z_sim_arr,wavelength_arr,"HI")
        inp_dict_tau.update({"wavelength_arr":wavelength_arr,"tau_total_HI_arr":tau_total_HI_arr})

    if flag_absorber.upper() == "METAL":
        tau_total_metal_arr = compute_optical_depth_for_given_specie_type(inp_dict,bool_arr_z_em,z_sim_arr,wavelength_arr,"METAL")
        inp_dict_tau.update({"wavelength_arr":wavelength_arr,"tau_total_metal_arr":tau_total_metal_arr})

    if flag_absorber.upper() == "DLA":
        tau_total_DLA_arr = compute_optical_depth_for_DLA(inp_dict,bool_arr_z_em,z_sim_arr,wavelength_arr)
        inp_dict_tau.update({"wavelength_arr":wavelength_arr,"tau_total_DLA_arr":tau_total_DLA_arr})

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

    return inp_dict


def compute_optical_depth_for_given_specie_type(inp_dict,bool_arr_z_em,z_sim_arr,wavelength_arr,specie):
    line_dict              = inp_dict["line_dict"]
    inp_path               = inp_dict["out_path"]

    if specie == "HI":
        inp_dict_tau           = inp_dict["inp_dict_tau_HI"]
        specie_trans_list_extr = inp_dict["HI_trans_list"]
    else:
        inp_dict_tau           = inp_dict["inp_dict_tau_metal"]
        specie_trans_list_extr = inp_dict["metal_trans_list"]

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

    for specie_trans in specie_trans_list_extr:
        tau_XI_arr              = inp_dict_tau["tau_" + specie_trans + "_arr"][bool_arr_z_em]
        lambda_rest_XI          = line_dict[specie_trans]["lambda_rest"]
        lambda_arr_z            = lambda_rest_XI * (1.0 + z_sim_arr)
        bool_arr                = (lambda_arr_z[0] <= wavelength_arr) & (wavelength_arr <= lambda_arr_z[-1])
        intrp_funct             = interp1d(lambda_arr_z,tau_XI_arr)
        tau_total_arr[bool_arr] = tau_total_arr[bool_arr] + intrp_funct(wavelength_arr[bool_arr])

    return tau_total_arr

def compute_optical_depth_for_DLA(inp_dict,bool_arr_z_em,z_sim_arr,wavelength_arr):
    line_dict      = inp_dict["line_dict"]
    inp_dict_tau   = inp_dict["inp_dict_tau_DLA"]
    tau_total_arr  = np.zeros(wavelength_arr.shape[0])
    tau_DLA_arr    = inp_dict_tau["tau_DLA_arr"]
    lambda_arr_z   = inp_dict_tau["lambda_arr_DLA"]
    bool_arr       = (lambda_arr_z[0] <= wavelength_arr) & (wavelength_arr <= lambda_arr_z[-1])
    intrp_funct    = interp1d(lambda_arr_z,tau_DLA_arr)
    tau_total_arr[bool_arr] = tau_total_arr[bool_arr] + intrp_funct(wavelength_arr[bool_arr])
    return tau_total_arr
