#!/usr/bin/env python3
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.

import numpy as np
from scipy.integrate import quad
from scipy.stats import expon
from scipy.special import spherical_jn, erf
from numpy import pi
from functools import lru_cache

# on retina screens, the default plots are way too small
# by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
# Qt5 will scale everything using the dpi in ~/.Xresources
import matplotlib
matplotlib.use("Qt5Agg")

# speed of light (exact)
SPEED_OF_LIGHT = 299792458 # m/s

# depth of the SNO detector (m)
# currently just converted 6800 feet -> meters
h = 2072.0

# radius of earth (m)
# from the "Earth Fact Sheet" at https://nssdc.gsfc.nasa.gov/planetary/factsheet/earthfact.html
# we use the volumetric mean radius here
R = 6.371e6

# Approximate dark matter velocity in m/s. The true distribution is expected to
# be a Maxwell Boltzmann distribution which is modulated annually by the
# earth's rotation around the sun, but we just assume a single constant
# velocity here. From Lewin and Smith Appendix B
DM_VELOCITY = 230e3

# Approximate earth velocity in the galactic rest frame (?)
# from Lewin and Smith Equation 3.6
EARTH_VELOCITY = 244e3

# Approximate dark matter density in GeV/m^3. From Tom Caldwell's thesis.
# from Lewin and Smith page 91
DM_DENSITY = 0.4e6

# Number density of scatterers in the Earth.
#
# FIXME: Currently just set to the number density of atoms in water. Need to
# update this for rock, and in fact this will change near the detector since
# there is water outside the AV.
DENSITY_WATER = 1e3 # In kg/m^3

# mean density of norite rock which is the rock surrounding SNO
# probably conservative
NORITE_DENSITY = 3e3 # In kg/m^3

# atomic masses for various elements
# from https://www.angelo.edu/faculty/kboudrea/periodic/structure_mass.htm
element_mass = {
  'H':1.0,
  'C':12.011,
  'O':15.9994,
  'Na':22.98977,
  'Mg':24.305,
  'Al':26.98154,
  'Si':28.0855,
  'K':39.0983,
  'Ca':40.08,
  'Mn':54.9380,
  'Fe':55.847,
  'Ti':47.90
}

# composition and mean density of norite rock which is the rock surrounding SNO
# the composition is from Table 3.2 in the SNOLAB User's handbook
water = {'composition':
 {'H':20,
  'O':80},
  'density':1e3}

# composition and mean density of the mantle
# from www.knowledgedoor.com/2/elements_handbook/element_abundances_in_the_earth_s_mantle.html
# density from hyperphysics.phy-astr.gsu.edu/hbase/Geophys/earthstruct.html
mantle = {'composition':
 {'O':44.33,
  'Mg':22.17,
  'Si':21.22,
  'Fe':6.3},
  'density':4.4e3}

# composition and mean density of norite rock which is the rock surrounding SNO
# the composition is from Table 3.2 in the SNOLAB User's handbook
norite = {'composition':
 {'H':0.15,
  'C':0.04,
  'O':46.0,
  'Na':2.2,
  'Mg':3.3,
  'Al':9.0,
  'Si':26.2,
  'K': 1.2,
  'Ca':5.2,
  'Mn':0.1,
  'Fe':6.2,
  'Ti':0.5},
  'density':3e3}

# Fiducial volume (m)
FIDUCIAL_RADIUS = 5

# Fiducial volume (m^3)
FIDUCIAL_VOLUME = 4*pi*FIDUCIAL_RADIUS**3/3

# proton mass from the PDG (2018)
PROTON_MASS = 0.938 # GeV

# proton mass from the PDG (2018)
ATOMIC_MASS_UNIT = 0.931 # GeV

# mass of Xenon in atomic units
XENON_MASS = 131.293

# mass of Neon in atomic units
NEON_MASS = 20.18

# mass of argon in atomic units
ARGON_MASS = 39.948

# mass of germanium in atomic units
GERMANIUM_MASS = 72.64

# mass of tungsten in atomic units
TUNGSTEN_MASS = 183.84

# mass of oxygen in atomic units
OXYGEN_MASS = 15.999

# mass of silicon in atomic units
SILICON_MASS = 28.0855

# mass of iron in atomic units
IRON_MASS = 55.845

# mass of magnesium in atomic units
MAGNESIUM_MASS = 24.305

# galactic escape velocity (m/s)
# from Tom Caldwell's thesis page 25
ESCAPE_VELOCITY = 244e3

# conversion constant from PDG
HBARC = 197.326978812e-3 # GeV fm

# Avogadros number (kg^-1)
N0 = 6.02214085774e26

def get_probability(r, l):
    """
    Returns the probability of a dark photon decaying in the SNO detector from
    a dark matter decay distributed uniformly in the Earth. Assumes that the
    depth of SNO is much larger than the dimensions of the SNO detector.
    """
    if r <= h:
        theta_min = 0
    elif r <= 2*R - h:
        theta_min = pi - np.arccos((h**2 + r**2 - 2*R*h)/(2*r*(R-h)))
    else:
        return 0
    return (1 + np.cos(theta_min))*np.exp(-r/l)/(2*l)

@lru_cache()
def get_probability2(l):
    p, err = quad(get_probability, 0, min(2*R-h,10*l), args=(l,), epsabs=0, epsrel=1e-7, limit=1000)

    return p

def get_event_rate(m, cs0, l, A):
    """
    Returns the event rate of leptons produced from self-destructing dark
    matter in the SNO detector for a given dark matter mass m, a cross section
    cs, and a mediator decay length l.
    """
    # For now we assume the event rate is constant throughout the earth, so we
    # are implicitly assuming that the cross section is pretty small.
    flux = DM_VELOCITY*DM_DENSITY/m

    #p, err = quad(get_probability, 0, min(2*R-h,10*l), args=(l,), epsabs=0, epsrel=1e-7, limit=1000)
    p = get_probability2(l)

    cs = get_cross_section(cs0, m, A)

    # FIXME: factor of 2 because the DM particle decays into two mediators?
    return p*cs*(N0/A)*flux*FIDUCIAL_VOLUME

def get_event_rate_sno(m, cs0, l, composition):
    """
    Returns the event rate of leptons produced from self-destructing dark
    matter in the SNO detector for a given dark matter mass m, a cross section
    cs, and a mediator decay length l.
    """
    rate = 0.0

    for element, mass_fraction in composition['composition'].items():
        rate += mass_fraction/100*get_event_rate(m,cs0,l,element_mass[element])

    return rate*composition['density']

def get_nuclear_form_factor(A, e):
    """
    Returns the nuclear form factor for a WIMP-nucleus interaction.

    From Tom Caldwell's thesis page 24.

    Also used Mark Pepin's thesis page 50
    """
    # mass of xenon nucleus
    mn = A*ATOMIC_MASS_UNIT

    # calculate approximate size of radius
    s = 0.9 # fm
    a = 0.52 # fm
    c = 1.23*A**(1/3) - 0.60 # fm
    r1 = np.sqrt(c**2 + (7/3)*pi**2*a**2 - 5*s**2)

    q = np.sqrt(2*mn*e)

    if q*r1/HBARC < 1e-10:
        return 1.0

    # Helm form factor
    # from Mark Pepin's thesis page 50
    f = 3*spherical_jn(1,q*r1/HBARC)*np.exp(-(q*s/HBARC)**2/2)/(q*r1/HBARC)

    return f

def get_cross_section(cs0, m, A):
    """
    Returns the WIMP cross section from the target-independent WIMP-nucleon
    cross section cs0 at zero momentum transfer.

    From Tom Caldwell's thesis page 21.
    """
    # mass of xenon nucleus
    mn = A*ATOMIC_MASS_UNIT

    # reduced mass of the nucleus and the WIMP
    mr = (m*mn)/(m + mn)

    # reduced mass of the proton and the WIMP
    mp = (m*PROTON_MASS)/(m + PROTON_MASS)

    return cs0*A**2*mr**2/mp**2

def get_differential_event_rate_xenon(e, m, cs0, A):
    """
    Returns the event rate of WIMP scattering in the Xenon 100T detector.
    """
    # mass of nucleus
    mn = A*ATOMIC_MASS_UNIT

    # reduced mass of the nucleus and the WIMP
    mr = (m*mn)/(m + mn)

    cs = get_cross_section(cs0, m, A)

    v0 = DM_VELOCITY

    vesc = ESCAPE_VELOCITY

    # earth's velocity through the galaxy
    ve = EARTH_VELOCITY

    # minimum wimp velocity needed to produce a recoil of energy e
    vmin = np.sqrt(mn*e/2)*(mn+m)*SPEED_OF_LIGHT/(mn*m)

    f = get_nuclear_form_factor(A, e)

    x = vmin/v0
    y = ve/v0
    z = vesc/v0

    # Equation 3.49 in Mark Pepin's thesis
    k0 = (pi*v0**2)**(3/2)

    # Equation 3.49 in Mark Pepin's thesis
    k1 = k0*(erf(z)-(2/np.sqrt(pi))*z*np.exp(-z**2))

    # From Mark Pepin's CDMS thesis page 59
    if x <= z - y:
        I = (k0/k1)*(DM_DENSITY/(2*ve))*(erf(x+y) - erf(x-y) - (4/np.sqrt(pi))*y*np.exp(-z**2))
    elif x <= y + z:
        I = (k0/k1)*(DM_DENSITY/(2*ve))*(erf(z) - erf(x-y) - (2/np.sqrt(pi))*(y+z-x)*np.exp(-z**2))
    else:
        return 0

    return (N0/A)*(mn/(2*m*mr**2))*cs*f**2*I*SPEED_OF_LIGHT**2

def get_event_rate_xenon(m, cs, A, threshold):
    """
    Returns the event rate of WIMP scattering in a dark matter detector using
    an element with atomic mass A. Rate is in Hz kg^-1.
    """
    # mass of nucleus
    mn = A*ATOMIC_MASS_UNIT

    # reduced mass of the nucleus and the WIMP
    mr = (m*mn)/(m + mn)

    vesc = ESCAPE_VELOCITY

    # earth's velocity through the galaxy
    ve = EARTH_VELOCITY

    # max recoil they looked for was 40 keV
    emax = 2*mr**2*((ve+vesc)/SPEED_OF_LIGHT)**2/mn

    rate, err = quad(get_differential_event_rate_xenon, threshold, emax, args=(m,cs,A), epsabs=0, epsrel=1e-2, limit=1000)

    return rate

if __name__ == '__main__':
    import matplotlib.pyplot as plt

    ls = np.logspace(-1,8,1000)

    cs = 1e-50 # cm^2

    # FIXME: should use water density for L < 1 m, silicon density for L ~ 1
    # km, and iron for L >> 1 km
    rate = np.array([get_event_rate_sno(1, cs*1e-4, l, water) for l in ls])

    colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

    plt.figure()
    plt.subplot(111)
    plt.plot(ls, rate/np.max(rate),color=colors[0])
    plt.xlabel("Mediator Decay Length (m)")
    plt.ylabel("Event Rate (arbitrary units)")
    plt.axvline(x=FIDUCIAL_RADIUS, color=colors[1], ls='--', label="SNO radius")
    plt.axvline(x=h, ls='--', color=colors[2], label="Depth of SNO")
    plt.axvline(x=R, ls='--', color=colors[3], label="Earth radius")
    plt.gca().set_xscale('log')
    plt.gca().set_xlim((ls[0], ls[-1]))
    plt.title("Event Rate of Self-Destructing Dark Matter in SNO")
    plt.legend()

    # threshold is ~5 keVr
    xenon_100t_threshold = 5e-6

    # threshold is ~1 keVr
    cdms_threshold = 1e-6

    # threshold is ~100 eV
    # FIXME: is this correct?
    cresst_threshold = 1e-7

    ms = np.logspace(-2,3,200)
    cs0s = np.logspace(-50,-40,200)

    mm, cs0cs0 = np.meshgrid(ms, cs0s)

    rate1 = np.empty(mm.shape)
    rate2 = np.empty(mm.shape)
    rate3 = np.empty(mm.shape)
    rate4 = np.empty(mm.shape)
    rate5 = np.empty(mm.shape)
    rate6 = np.empty(mm.shape)
    for i in range(mm.shape[0]):
        print("\r%i/%i" % (i+1,mm.shape[0]),end='')
        for j in range(mm.shape[1]):
            rate1[i,j] = get_event_rate_xenon(mm[i,j], cs0cs0[i,j]*1e-4, XENON_MASS, xenon_100t_threshold)
            rate2[i,j] = get_event_rate_sno(mm[i,j], cs0cs0[i,j]*1e-4, 1.0, water)
            rate3[i,j] = get_event_rate_xenon(mm[i,j], cs0cs0[i,j]*1e-4, GERMANIUM_MASS, cdms_threshold)
            rate4[i,j] = get_event_rate_xenon(mm[i,j], cs0cs0[i,j]*1e-4, TUNGSTEN_MASS, cresst_threshold)
            rate5[i,j] = get_event_rate_sno(mm[i,j], cs0cs0[i,j]*1e-4, 1e3, norite)
            rate6[i,j] = get_event_rate_sno(mm[i,j], cs0cs0[i,j]*1e-4, 1e6, mantle)
    print()

    # Fiducial volume of the Xenon1T detector is 1042 +/- 12 kg
    # from arxiv:1705.06655
    xenon_100t_fiducial_volume = 1042 # kg

    # Livetime of the Xenon1T results
    # from arxiv:1705.06655
    xenon_100t_livetime = 34.2 # days

    plt.figure()
    plt.subplot(111)
    plt.gca().set_xscale('log')
    plt.gca().set_yscale('log')
    CS1 = plt.contour(mm,cs0cs0,rate1*3600*24*xenon_100t_fiducial_volume*xenon_100t_livetime,[10.0], colors=[colors[0]])
    CS2 = plt.contour(mm,cs0cs0,rate2*3600*24*668.8,[10.0],colors=[colors[1]])
    CS3 = plt.contour(mm,cs0cs0,rate3*3600*24*70.10,[10.0],colors=[colors[2]])
    # FIXME: I used 2.39 kg day because that's what CRESST-3 reports in their paper
    # but! I only use Tungsten here so do I need to multiply by the mass fraction of tungsten?
    CS4 = plt.contour(mm,cs0cs0,rate4*3600*24*2.39,[10.0],colors=[colors[3]])
    CS5 = plt.contour(mm,cs0cs0,rate5*3600*24*668.8,[10.0],colors=[colors[1]], linestyles=['dashed'])
    CS6 = plt.contour(mm,cs0cs0,rate6*3600*24*668.8,[10.0],colors=[colors[1]], linestyles=['dotted'])
    plt.clabel(CS1, inline=1, fmt="XENON1T", fontsize=10, use_clabeltext=True)
    plt.clabel(CS2, inline=1, fmt=r"SNO ($\mathrm{L}_V$ = 1 m)", fontsize=10)
    plt.clabel(CS3, inline=1, fmt="CDMSLite", fontsize=10)
    plt.clabel(CS4, inline=1, fmt="CRESST-3", fontsize=10)
    plt.clabel(CS5, inline=1, fmt=r"SNO ($\mathrm{L}_V$ = 1 km)", fontsize=10)
    plt.clabel(CS6, inline=1, fmt=r"SNO ($\mathrm{L}_V$ = 1000 km)", fontsize=10)
    plt.xlabel(r"$m_\chi$ (GeV)")
    plt.ylabel(r"WIMP-nucleon scattering cross section ($\mathrm{cm}^2$)")
    plt.title("Self-Destructing Dark Matter Limits")

    x = np.linspace(0,300e-6,1000)
    
    # reproducing Figure 2.1 in Tom Caldwell's thesis
    plt.figure()
    plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(XENON_MASS, x)**2,x)), label="Xe")
    plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(NEON_MASS, x)**2,x)), label="Ne")
    plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(ARGON_MASS, x)**2,x)), label="Ar")
    plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(GERMANIUM_MASS, x)**2,x)), label="Ge")
    plt.xlabel("Recoil Energy (keV)")
    plt.ylabel(r"$F^2(E)$")
    plt.gca().set_yscale('log')
    plt.legend()
    plt.gca().set_ylim((1e-5,1))
    plt.gca().set_xlim((0,300))
    plt.title("Nuclear Form Factors for various elements")
    plt.show()