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


pl.clf()
vmin,delta_eta1=WD.streamed_halo_function(yearly_modulation=True)
# delta_eta1 for yearly modulation calculated setting yearly_modulation=True
#
# defines c_tau (in GeV^-2) in terms of the reference cross section
# sigma_ref=|c_vec|^2*mu_chi_n^2/pi (in cm^2)
# where c_vec=[c4^1,c4^2,c5^1,c5^2,c6^1,c6^2]
# and mu_chi_n is the WIMP-nucleon reduced mass.
def c_tau(mchi,sigma_ref,c_tilde):
    hbarc2=0.389e-27
    mn=0.931
    mu=mchi*mn/(mchi+mn)
    c_abs=(mu**2*hbarc2/(np.pi*sigma_ref))**(-0.5)
    return c_abs*c_tilde

# c_tilde=c_vec/|c_vec| is the unit vector pointing in the direction of c_vec
# that minimizes the tension with the DAMA modulation amplitudes
c_tilde=np.array([-0.00139183, -0.00145925,-0.00318353, -0.01660634, 0.69199713, 0.72169939])

# values of WIMP mass, mass splitting and sigma_ref that minimize the tension with DAMA
# modulation amplitudes
mchi=11.64102564102564
delta=23.73913043478261
sigma_ref=4.684124873205523e-28
# calculates c_vec in terms of sigma_ref, mchi,c_tilde
c_vec=c_tau(mchi,sigma_ref,c_tilde)

# defines a Hamiltonian that maps mchi, sigma_ref and c_tilde into the couplings c_4^tau, c_5^tau, c_6^tau

wc={4: lambda mchi,s,c: c_tau(mchi,s,c)[0:2],5: lambda mchi,s,c: c_tau(mchi,s,c)[2:4],6: lambda mchi,s,c: c_tau(mchi,s,c)[4:6]}
c_456=WD.eft_hamiltonian('c_456',wc)

# loads the response functions (j_chi=0.5 is default valu for spin)
WD.load_response_functions(WD.DAMA_LIBRA_2019,c_456,verbose=False)

# calculates the modulation amplitudes in terms of mchi, sigma_ref and c_tilde
e_vec,dama_th=WD.wimp_dd_rate(WD.DAMA_LIBRA_2019,c_456,vmin,delta_eta1,mchi,delta=delta,s=sigma_ref,c=c_tilde)
pl.plot(e_vec,dama_th,'-k',linewidth=1)


# modulation amplitudes measured by DAMA-LIBRA (from Fig. 11 of NUCL. PHYS. AT. ENERGY 19 (2018) 307-325)
experimental_s1=np.array([0.0232    , 0.0164    , 0.0178    , 0.019     , 0.0178,0.0109, 0.011, 0.004, 0.0065, 0.0066,0.0009,0.0016,0.0007,0.0016, 0.00039922])

errors_on_s1=np.array([0.0052, 0.0043, 0.0028    , 0.0029    , 0.0028,0.0025    , 0.0022    , 0.002     , 0.002, 0.0019,0.0018    , 0.0018    , 0.0018    , 0.0018    , 0.00046397])

errors_on_energy=np.array([0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25,0.25, 0.25, 0.25, 4.  ])
pl.errorbar(e_vec,experimental_s1,yerr=errors_on_s1,xerr=errors_on_energy,fmt='none',ecolor='red')

pl.xlabel('E$^{\prime}$ (keV)')
pl.ylabel('$S^{1}_{[E_1^{\prime},E_2^{\prime}]}$ (events/kg/day/keV)')

pl.show()