import numpy as np

import matplotlib.pyplot as plt

from common_lib.plot_setting import load_plt

from common_lib.parameter_files.generate_line_dict import generate_line_dict_main



def CoG_main():
    specie_trans  = "HI-1215"
    line_dict_all = generate_line_dict_main()
    line_dict     = line_dict_all[specie_trans]
    prop_dict     = generate_CoG_properties()

    lambda_rest   = line_dict["lambda_rest"]
    z_abs         = 0.0
    lambda_c      = lambda_rest * (1 + z_abs)
    wavelength    = np.linspace(-50.0,50.0,10001) + lambda_c
    fig,ax        = plt.subplots(1)
    fig.set_figwidth(11.0)
    fig.set_figheight(10.0)

    for b_indx in range(len(prop_dict)):
        prop_sub_dict = prop_dict[b_indx]
        COMPUTE_CoG_FOR_GIVEN_b_PARAMETER(ax,lambda_c,wavelength,line_dict,prop_sub_dict)

    SET_AX_PROPERTY(ax)
    ax.legend(loc="best")
    plt.savefig("COG.pdf",dpi=200)

def COMPUTE_CoG_FOR_GIVEN_b_PARAMETER(ax,lambda_c,wavelength,line_dict,prop_sub_dict):
    b_val         = prop_sub_dict["b_val"]
    color         = prop_sub_dict["color"]
    ls            = prop_sub_dict["ls"]
    lw            = prop_sub_dict["lw"]
    alpha         = prop_sub_dict["alpha"]
    marker        = prop_sub_dict["marker"]
    markersize    = prop_sub_dict["markersize"]

    N_HI_arr      = np.logspace(np.log10(1e12),np.log10(1e20),200)
    W_eq_voigt    = np.zeros(N_HI_arr.shape[0])

    for indx,N_HI in enumerate(N_HI_arr):
        Flux_voigt       = VOIGT_FUNCTION(wavelength,b_val,N_HI,lambda_c,line_dict)
        W_eq_voigt[indx] = EQUVIVALENT_WIDTH(wavelength,Flux_voigt)

    label = r"$b = " + str("%.0f"%b_val) + r"$ km s$^{-1}$"
    ax.plot(np.log10(N_HI_arr),np.log10(W_eq_voigt),lw=lw,ls=ls,color=color,alpha=alpha)
    ax.plot(np.log10(N_HI_arr)[::10],np.log10(W_eq_voigt)[::10],ls="",color=color,label=label,marker=marker,markersize=markersize,alpha=alpha)

def SET_AX_PROPERTY(ax):
    plt.xlabel(r"$\log({\rm N_{HI}} / {\rm cm}^{-2})$",size=20)
    plt.ylabel(r"$\log(W \; [\AA])$",size=20)



def VOIGT_FUNCTION(wavelength,b,N_HI,lambda_c,line_dict):
	def voigt(x, y):
		from scipy.special import wofz # Faddeeva function
		z = x + 1j*y
		I = wofz(z).real
		return I

	I_alpha		=	line_dict["i_alpha"]
	damp_gamma	= 	line_dict["damp_gamma"]
	light_speed	=	3e5
	B			=	light_speed * (wavelength - lambda_c)  / (wavelength * b)
	A			=	lambda_c * damp_gamma * 1e-13 / (4 * np.pi * b)
	tau_c		=	light_speed * N_HI * I_alpha / (b * np.pi**0.5)
	tau			=	tau_c  * voigt(B,A) # (np.exp(-B**2.0) + (0.564*A / (A**2.0+B**2.0))) #voigt(B,A)
	return np.exp(-1*tau)


def generate_CoG_properties():
    prop_dict = {0:{"b_val":5.0, "ls":"--","lw":2.50,"color":"blue", "marker":"o","markersize":9, "alpha":0.65},\
                 1:{"b_val":10.0,"ls":"-", "lw":1.75,"color":"red",  "marker":"*","markersize":15,"alpha":0.65},\
                 2:{"b_val":20.0,"ls":"-.","lw":3.50,"color":"black","marker":"D","markersize":9, "alpha":0.65},\
                 }
    return prop_dict

def EQUVIVALENT_WIDTH(wavelength,Flux):
	from scipy.integrate import simps
	return simps(1.0-Flux,wavelength)


CoG_main()







