import numpy as np
import matplotlib.pyplot as plt

"""
Script originally written by Ben Hokanson-Fasig
"""


########################## Generate the data to plot ##########################
energies = 10**np.arange(15.5, 20.5, 0.5)
prob_all = [0.026102207436026353, 0.22207215965381144, 0.5216687959104875,
            0.6475666483333804, 0.6877039526636356, 0.773031394004548,
            0.792320790547285, 0.7798359568015701, 0.820271800483975,
            0.8797223073849653]
err_all = [0.018352019249675763, 0.04842096986969432, 0.053830783088298126,
           0.033859008395343816, 0.03547864548770085, 0.03137797157385291,
           0.029827650169126828, 0.022430504846313797, 0.01969317261724343,
           0.01655925303773214]
prob_surf = [0.0, 0.03167443388226018, 0.18712332645067217,
             0.3084496990399548, 0.31922577587441525, 0.3706411871325141,
             0.3695828389965637, 0.35486278003830346, 0.37170091959491947,
             0.36833695465704597]
err_surf = [0.0, 0.018286932286523422, 0.03224016326092905,
            0.023368141509953675, 0.024172148086513412, 0.02172716880345247,
            0.020371565914620162, 0.015131000505149506, 0.01325664464024523,
            0.01071496080564502]

prob_deep = [0.026102207436026353, 0.20096329396035684, 0.46587916328503953,
             0.5696250473880095, 0.6031702064894984, 0.7034860944055573,
             0.7228583323370091, 0.6968159248146896, 0.7368183949754757,
             0.8227099224520789]
err_deep = [0.018352019249675763, 0.046062212689570316, 0.050870957394092205,
            0.031756054408821704, 0.03322662357500077, 0.029933262762294754,
            0.028490174587762356, 0.021202958053604764, 0.018664528645988655,
            0.016013685974963686]



################# Define a function to plot with error region #################
def plot_with_error(x, y, yerr, color=None, alpha=0.25, label=None):
    """
    Plot data points with a shaded error region.

    Plots the given data points with a logarithmic x scale and produces a
    shaded region around the data with a size specified by `yerr`.

    Parameters
    ----------
    x : array_like
        The horizontal coordinates of the data points.
    y : array_like
        The vertical coordinates of the data points.
    yerr : array_like
        The errors on the values of `y`.
    color : str, optional
        The color of the line and shaded region.
    alpha : float, optional
        The opacity of the shaded region.
    label : str, optional
        The legend label for the line.
    """
    # Plot the data with solid lines between dots on a logarithmic x scale
    plt.semilogx(x, y, '.-', color=color, label=label)
    # Fill the error region with a semi-transparent color fill
    plt.fill_between(x, [p-e for p, e in zip(y, yerr)],
                     [p+e for p, e in zip(y, yerr)],
                     color=color, alpha=0.25)



#################### Plot surface coincidence probabilities ###################
# Set the size of the plot
plt.figure(figsize=(7, 5))
# Plot the data using our function defined above
plot_with_error(energies, prob_all, err_all, color='C0',
                label='Any 3 Antennas > 3$\\sigma$')
plot_with_error(energies, prob_surf, err_surf, color='C1',
                label='3 Surface Antennas > 3$\\sigma$')
plot_with_error(energies, prob_deep, err_deep, color='C2',
                label='3 Deep Antennas > 3$\\sigma$')
# Add a grid of dotted lines
plt.grid(ls=':')
# Give the plot a title and axis labels
plt.title("Probability for 3$\\sigma$ Signals in RNO Station", fontsize=14)
plt.xlabel("Neutrino Energy [eV]", fontsize=12)
plt.ylabel("Probability of Antenna Coincidence", fontsize=12)
# Add a legend for the given line labels
plt.legend()
# Set the fontsize of all axis labels
for item in (plt.gca().get_xticklabels() + plt.gca().get_yticklabels()):
    item.set_fontsize(12)
# Limit the axis ranges in x and y
plt.xlim(1e16, 1e20)
plt.ylim(0, 1)
# Save the plot as a PDF so it scales nicely in presentations
plt.savefig('coincidence_rates.pdf')