#!/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/>.

from __future__ import print_function, division
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 collections import namedtuple
from functools import update_wrapper
from threading import RLock

# Backport of lru_cache from http://code.activestate.com/recipes/578078-py26-and-py30-backport-of-python-33s-lru-cache/
_CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])

class _HashedSeq(list):
    __slots__ = 'hashvalue'

    def __init__(self, tup, hash=hash):
        self[:] = tup
        self.hashvalue = hash(tup)

    def __hash__(self):
        return self.hashvalue

def _make_key(args, kwds, typed,
             kwd_mark = (object(),),
             fasttypes = {int, str, frozenset, type(None)},
             sorted=sorted, tuple=tuple, type=type, len=len):
    'Make a cache key from optionally typed positional and keyword arguments'
    key = args
    if kwds:
        sorted_items = sorted(kwds.items())
        key += kwd_mark
        for item in sorted_items:
            key += item
    if typed:
        key += tuple(type(v) for v in args)
        if kwds:
            key += tuple(type(v) for k, v in sorted_items)
    elif len(key) == 1 and type(key[0]) in fasttypes:
        return key[0]
    return _HashedSeq(key)

def lru_cache(maxsize=100, typed=False):
    """Least-recently-used cache decorator.

    If *maxsize* is set to None, the LRU features are disabled and the cache
    can grow without bound.

    If *typed* is True, arguments of different types will be cached separately.
    For example, f(3.0) and f(3) will be treated as distinct calls with
    distinct results.

    Arguments to the cached function must be hashable.

    View the cache statistics named tuple (hits, misses, maxsize, currsize) with
    f.cache_info().  Clear the cache and statistics with f.cache_clear().
    Access the underlying function with f.__wrapped__.

    See:  http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used

    """

    # Users should only access the lru_cache through its public API:
    #       cache_info, cache_clear, and f.__wrapped__
    # The internals of the lru_cache are encapsulated for thread safety and
    # to allow the implementation to change (including a possible C version).

    def decorating_function(user_function):

        cache = dict()
        stats = [0, 0]                  # make statistics updateable non-locally
        HITS, MISSES = 0, 1             # names for the stats fields
        make_key = _make_key
        cache_get = cache.get           # bound method to lookup key or return None
        _len = len                      # localize the global len() function
        lock = RLock()                  # because linkedlist updates aren't threadsafe
        root = []                       # root of the circular doubly linked list
        root[:] = [root, root, None, None]      # initialize by pointing to self
        nonlocal_root = [root]                  # make updateable non-locally
        PREV, NEXT, KEY, RESULT = 0, 1, 2, 3    # names for the link fields

        if maxsize == 0:

            def wrapper(*args, **kwds):
                # no caching, just do a statistics update after a successful call
                result = user_function(*args, **kwds)
                stats[MISSES] += 1
                return result

        elif maxsize is None:

            def wrapper(*args, **kwds):
                # simple caching without ordering or size limit
                key = make_key(args, kwds, typed)
                result = cache_get(key, root)   # root used here as a unique not-found sentinel
                if result is not root:
                    stats[HITS] += 1
                    return result
                result = user_function(*args, **kwds)
                cache[key] = result
                stats[MISSES] += 1
                return result

        else:

            def wrapper(*args, **kwds):
                # size limited caching that tracks accesses by recency
                key = make_key(args, kwds, typed) if kwds or typed else args
                with lock:
                    link = cache_get(key)
                    if link is not None:
                        # record recent use of the key by moving it to the front of the list
                        root, = nonlocal_root
                        link_prev, link_next, key, result = link
                        link_prev[NEXT] = link_next
                        link_next[PREV] = link_prev
                        last = root[PREV]
                        last[NEXT] = root[PREV] = link
                        link[PREV] = last
                        link[NEXT] = root
                        stats[HITS] += 1
                        return result
                result = user_function(*args, **kwds)
                with lock:
                    root, = nonlocal_root
                    if key in cache:
                        # getting here means that this same key was added to the
                        # cache while the lock was released.  since the link
                        # update is already done, we need only return the
                        # computed result and update the count of misses.
                        pass
                    elif _len(cache) >= maxsize:
                        # use the old root to store the new key and result
                        oldroot = root
                        oldroot[KEY] = key
                        oldroot[RESULT] = result
                        # empty the oldest link and make it the new root
                        root = nonlocal_root[0] = oldroot[NEXT]
                        oldkey = root[KEY]
                        oldvalue = root[RESULT]
                        root[KEY] = root[RESULT] = None
                        # now update the cache dictionary for the new links
                        del cache[oldkey]
                        cache[key] = oldroot
                    else:
                        # put result in a new link at the front of the list
                        last = root[PREV]
                        link = [last, root, key, result]
                        last[NEXT] = root[PREV] = cache[key] = link
                    stats[MISSES] += 1
                return result

        def cache_info():
            """Report cache statistics"""
            with lock:
                return _CacheInfo(stats[HITS], stats[MISSES], maxsize, len(cache))

        def cache_clear():
            """Clear the cache and cache statistics"""
            with lock:
                cache.clear()
                root = nonlocal_root[0]
                root[:] = [root, root, None, None]
                stats[:] = [0, 0]

        wrapper.__wrapped__ = user_function
        wrapper.cache_info = cache_info
        wrapper.cache_clear = cache_clear
        return update_wrapper(wrapper, user_function)

    return decorating_function

# 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

# Taken from https://raw.githubusercontent.com/mwaskom/seaborn/c73055b2a9d9830c6fbbace07127c370389d04dd/seaborn/utils.py
def despine(fig=None, ax=None, top=True, right=True, left=False,
            bottom=False, offset=None, trim=False):
    """Remove the top and right spines from plot(s).

    fig : matplotlib figure, optional
        Figure to despine all axes of, default uses current figure.
    ax : matplotlib axes, optional
        Specific axes object to despine.
    top, right, left, bottom : boolean, optional
        If True, remove that spine.
    offset : int or dict, optional
        Absolute distance, in points, spines should be moved away
        from the axes (negative values move spines inward). A single value
        applies to all spines; a dict can be used to set offset values per
        side.
    trim : bool, optional
        If True, limit spines to the smallest and largest major tick
        on each non-despined axis.

    Returns
    -------
    None

    """
    # Get references to the axes we want
    if fig is None and ax is None:
        axes = plt.gcf().axes
    elif fig is not None:
        axes = fig.axes
    elif ax is not None:
        axes = [ax]

    for ax_i in axes:
        for side in ["top", "right", "left", "bottom"]:
            # Toggle the spine objects
            is_visible = not locals()[side]
            ax_i.spines[side].set_visible(is_visible)
            if offset is not None and is_visible:
                try:
                    val = offset.get(side, 0)
                except AttributeError:
                    val = offset
                _set_spine_position(ax_i.spines[side], ('outward', val))

        # Potentially move the ticks
        if left and not right:
            maj_on = any(
                t.tick1line.get_visible()
                for t in ax_i.yaxis.majorTicks
            )
            min_on = any(
                t.tick1line.get_visible()
                for t in ax_i.yaxis.minorTicks
            )
            ax_i.yaxis.set_ticks_position("right")
            for t in ax_i.yaxis.majorTicks:
                t.tick2line.set_visible(maj_on)
            for t in ax_i.yaxis.minorTicks:
                t.tick2line.set_visible(min_on)

        if bottom and not top:
            maj_on = any(
                t.tick1line.get_visible()
                for t in ax_i.xaxis.majorTicks
            )
            min_on = any(
                t.tick1line.get_visible()
                for t in ax_i.xaxis.minorTicks
            )
            ax_i.xaxis.set_ticks_position("top")
            for t in ax_i.xaxis.majorTicks:
                t.tick2line.set_visible(maj_on)
            for t in ax_i.xaxis.minorTicks:
                t.tick2line.set_visible(min_on)

        if trim:
            # clip off the parts of the spines that extend past major ticks
            xticks = ax_i.get_xticks()
            if xticks.size:
                firsttick = np.compress(xticks >= min(ax_i.get_xlim()),
                                        xticks)[0]
                lasttick = np.compress(xticks <= max(ax_i.get_xlim()),
                                       xticks)[-1]
                ax_i.spines['bottom'].set_bounds(firsttick, lasttick)
                ax_i.spines['top'].set_bounds(firsttick, lasttick)
                newticks = xticks.compress(xticks <= lasttick)
                newticks = newticks.compress(newticks >= firsttick)
                ax_i.set_xticks(newticks)

            yticks = ax_i.get_yticks()
            if yticks.size:
                firsttick = np.compress(yticks >= min(ax_i.get_ylim()),
                                        yticks)[0]
                lasttick = np.compress(yticks <= max(ax_i.get_ylim()),
                                       yticks)[-1]
                ax_i.spines['left'].set_bounds(firsttick, lasttick)
                ax_i.spines['right'].set_bounds(firsttick, lasttick)
                newticks = yticks.compress(yticks <= lasttick)
                newticks = newticks.compress(newticks >= firsttick)
                ax_i.set_yticks(newticks)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser("plot fit results")
    parser.add_argument("--save", action="store_true", default=False, help="save plots")
    args = parser.parse_args()

    if args.save:
        # default \textwidth for a fullpage article in Latex is 16.50764 cm.
        # You can figure this out by compiling the following TeX document:
        #
        #    \documentclass{article}
        #    \usepackage{fullpage}
        #    \usepackage{layouts}
        #    \begin{document}
        #    textwidth in cm: \printinunitsof{cm}\prntlen{\textwidth}
        #    \end{document}

        width = 16.50764
        width /= 2.54 # cm -> inches
        # According to this page:
        # http://www-personal.umich.edu/~jpboyd/eng403_chap2_tuftegospel.pdf,
        # Tufte suggests an aspect ratio of 1.5 - 1.6.
        height = width/1.5
        FIGSIZE = (width,height)

        import matplotlib.pyplot as plt

        font = {'family':'serif', 'serif': ['computer modern roman']}
        plt.rc('font',**font)

        plt.rc('text', usetex=True)
    else:
        # 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")

        import matplotlib.pyplot as plt

        # Default figure size. Currently set to my monitor width and height so that
        # things are properly formatted
        FIGSIZE = (13.78,7.48)

        # Make the defalt font bigger
        plt.rc('font', size=22)

    plt.rcParams['figure.figsize'] = FIGSIZE

    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']

    fig = plt.figure(1)
    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.legend(framealpha=1.0)
    despine(fig,trim=True)
    plt.tight_layout()

    # 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

    fig = plt.figure(2)
    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.tight_layout()

    x = np.linspace(0,300e-6,1000)
    
    # reproducing Figure 2.1 in Tom Caldwell's thesis
    fig = plt.figure(3)
    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.tight_layout()

    if args.save:
        fig = plt.figure(1)
        plt.savefig("sddm_event_rate.pdf")
        plt.savefig("sddm_event_rate.eps")
        fig = plt.figure(2)
        plt.savefig("sno_sensitivity.pdf")
        plt.savefig("sno_sensitivity.eps")
        fig = plt.figure(3)
        plt.savefig("nuclear_form_factors.pdf")
        plt.savefig("nuclear_form_factors.eps")
    else:
        plt.figure(1)
        plt.title("Event Rate of Self-Destructing Dark Matter in SNO")
        plt.figure(2)
        plt.title("Self-Destructing Dark Matter Limits")
        plt.figure(3)
        plt.title("Nuclear Form Factors for various elements")

        plt.show()