from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
from math import log10
from scipy.special import wofz # Faddeeva function

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.echo_functions import \
process_completion_print_function

def method_1_add_DLA_using_column_density_main(inp_dict):

    from WEAVEify_routines.step_02_common_routines.tau_summation_routines import \
    tau_summation_routines_main


    CHECK_AND_INITIALIZE_KEY_VALUE(inp_dict,"inp_dict_tau",{})

    inp_dict       = generate_DLA_property_dict(inp_dict)
    line_dict      = inp_dict["line_dict"]
    DLA_prop_dict  = inp_dict["DLA_prop_dict"]
    z_arr_sim      = inp_dict["z_arr"]
    N_DLA          = len(DLA_prop_dict)
    specie         = "HI"
    trans_list     = ["1215","1025"]
    lambda_HI_lya  = line_dict["HI-1215"]["lambda_rest"]
    lambda_arr_DLA = (1.0 + z_arr_sim) * lambda_HI_lya
    tau_DLA_arr    = np.zeros(lambda_arr_DLA.shape[0])

    for DLA_indx in range(N_DLA):
        DLA_prop_sub_dict = DLA_prop_dict[DLA_indx]
        log_N_XI          = DLA_prop_sub_dict["log_N_" + specie]
        z_XI              = DLA_prop_sub_dict["z_" + specie]
        b_XI              = DLA_prop_sub_dict["b_" + specie]

        flag_log_N_XI     = log_N_XI is not None
        flag_z_XI         = z_XI is not None
        flag_b_XI         = b_XI is not None

        if flag_log_N_XI & flag_z_XI & flag_b_XI:
            for trans in trans_list:
                line_sub_dict  = line_dict[specie + "-" + trans]
                lambda_rest    = line_sub_dict["lambda_rest"]
                lambda_c       = (1.0 + z_XI) * lambda_rest
                tau_arr_trans  = VOIGT_FUNCTION_FOR_DLA(lambda_arr_DLA,b_XI,log_N_XI,lambda_c,line_sub_dict)
                tau_DLA_arr    = tau_DLA_arr + tau_arr_trans

    process_completion_print_function(inp_dict,"Optical Depth calculation for DLA systems is completed")
    inp_dict_tau_DLA = {"tau_DLA_arr":tau_DLA_arr,"lambda_arr_DLA":lambda_arr_DLA}
    inp_dict.update({"inp_dict_tau_DLA":inp_dict_tau_DLA})

    inp_dict = tau_summation_routines_main(inp_dict,flag_absorber = "DLA")

    return inp_dict


def VOIGT_FUNCTION_FOR_DLA(wavelength,b,log_N_XI,lambda_c,line_dict):
    def voigt(x, y):
        z = x + 1j*y
        I = wofz(z).real
        return I

    I_alpha     = line_dict["i_alpha"]
    damp_gamma  = line_dict["damp_gamma"]
    lambda_rest = line_dict["lambda_rest"]
    N_XI        = 10**log_N_XI
    light_speed = 3e5
    B           = light_speed * (wavelength - lambda_c)  / (wavelength * b)
    A           = lambda_rest * damp_gamma * 1e-13 / (4 * np.pi * b)
    tau_c       = light_speed * N_XI * I_alpha / (b * np.pi**0.5)
    tau         = tau_c  * voigt(B,A)
    return tau

def generate_DLA_property_dict(inp_dict):
    log_N_HI_DLA = inp_dict["log_N_HI_DLA"]
    b_DLA        = inp_dict["b_DLA"]
    z_DLA        = inp_dict["z_DLA"]

    log_N_HI_DLA_list = convert_DLA_parameter_to_list_if_needed(log_N_HI_DLA)
    b_DLA_list        = convert_DLA_parameter_to_list_if_needed(b_DLA)
    z_DLA_list        = convert_DLA_parameter_to_list_if_needed(z_DLA)

    DLA_prop_dict     = {}
    N_DLA             = len(log_N_HI_DLA_list)

    for DLA_indx in range(N_DLA):
        b_DLA_val         = float(b_DLA_list[DLA_indx])
        z_DLA_val         = float(z_DLA_list[DLA_indx])
        log_N_HI_DLA_val  = float(log_N_HI_DLA_list[DLA_indx])
        DLA_prop_sub_dict = {DLA_indx:{"log_N_HI":log_N_HI_DLA_val,\
                                       "b_HI":b_DLA_val,\
                                       "z_HI":z_DLA_val}}
        DLA_prop_dict.update(DLA_prop_sub_dict)

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

    return inp_dict

def convert_DLA_parameter_to_list_if_needed(param_DLA_list):
    flag_list_param = type(param_DLA_list) is list
    if flag_list_param is False:
        if param_DLA_list is not None:
            param_DLA_list = [param_DLA_list]
        else:
            param_DLA_list = []
    return param_DLA_list




