import WimPyDD as WD
import numpy as np
import matplotlib.pyplot as pl


nai=WD.DAMA_LIBRA_2019.target
na,i=nai.element
#N.B. using the DAMA_LIBRA target sodium and iodine have the quenching attribute:
# 'quenching' in dir(na) -> True
# 'quenching' in dir(i) -> True
# the routine diff_rate interprets the argument er as electron-equivalent energy
# E_ee=Q(E_R)*E_R with Q the quenching and E_R the nuclear recoil energy.

# On the other hand:
# 'quenching' in dir(WD.Na) -> False
# 'quenching' in dir(WD.I) -> False
# i.e. the pre-defined targets do not have quenching information.
# if used in diff_rate the argument er is interpreted as nuclear recoil energy. 

#Define Hamiltonian in terms of two parameters: the effective scale M and r=cn/cp
# (ratio  of WIMP-neutron over WIMP-proton couplings) 
wc={1: lambda M, r=1 : [(1.+r)/M**2, (1-r)/M**2] }
SI_M=WD.eft_hamiltonian('SI_M',wc)

er_vec=np.linspace(0.1,20,100) 
mchi=20. 
vmin,delta_eta0=WD.streamed_halo_function(n_vmin_bin=1000)# Maxwellian with standard parameters
diff_rate=[WD.diff_rate(nai, SI_M, mchi, er, vmin, delta_eta0,M=1e3) for er in er_vec]
# N.B the r parameter is not passed, default value r=1 used. 
# N.B.2 also the argument exposure is not passed, so the rate is normalized to events/kg/day/keV. 
pl.plot(er_vec,diff_rate,'-b',label=r'$dR/dE_{ee}$, NaI')

# loads response functions (j_chi=0.5 spin default value)
# from the directory WimPyDD/Experiments/DAMA_LIBRA_2019/Response_functions/spin_1_2/
# since SI_M.coeff_squared_list -> [(1, 1)]
# the file c_1_c_1.npy is loaded.

WD.load_response_functions(WD.DAMA_LIBRA_2019,SI_M,verbose=False)
# before calling load_response_functions the dictionary attribute
# WD.DAMA_LIBRA_2019.response_functions is empty. After the call:
# [(hamiltonian.name,spin) for hamiltonian,spin in WD.DAMA_LIBRA_2019.response_functions.keys()] -> [('SI_M', 0.5)], i.e. the response functions for the SI_M hamiltonian and spin 1/2 are loaded.

e_prime,s0=WD.wimp_dd_rate(WD.DAMA_LIBRA_2019, SI_M, vmin, delta_eta0, mchi,M=1e3)
#N.B the output is in events/kg/day/keV because in the WimPyDD/Experiments/DAMAthe exposure is set
# to 1/delta_e, with delta_e the with of each bin:
#WD.DAMA_LIBRA_2019.exposure  -> [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.125]

pl.step(e_prime,s0,'-r',linewidth=2, label=r'$S^0_{[E^\prime_1,E^\prime_2]}$, NaI')

pl.legend(loc='upper right', ncol=2,frameon=False,fontsize=13)
pl.xlabel('E$^{\prime}$ or $E_{ee}$ (keV)')
pl.ylabel(' $dR/dE_R$ (events/kg/day/keV)') 

pl.yscale('log')

pl.show()