import numpy as np
import matplotlib.pyplot as plt

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

########################## Generate the data to plot ##########################
volumes_16p5 = np.array([0.009474834, 0.2366, 0.3711, 0.506, 0.6822])
volumes_18p0 = np.array([1.481634, 3.89, 7.013, 10.56, 19.99])
coin_16p5 = np.array([1.000000, 0.592732, 0.214104, 0.067033, 0.014277])
coin_18p0 = np.array([1.000000, 0.723583, 0.386608, 0.246627, 0.125687])
coin_vol_16p5 = coin_16p5 * volumes_16p5
coin_vol_18p0 = coin_18p0 * volumes_18p0



###################### Plot surface coincidence fractions #####################
xlabels = ["Surface", "15 m", "60 m", "100 m", "200 m"]
width = 0.4
sep = 0.2
n_bars = len(coin_16p5)
# Set the size of the plot
plt.figure(figsize=(8, 5))
# Set up horizontal dotted grid lines
plt.grid(axis='y', ls=':')
# Plot two sets of bar graphs, one for each data set
plt.bar(np.arange(n_bars)-sep, coin_16p5,
        width=width, alpha=0.8, label='1e16.5 eV')
plt.bar(np.arange(n_bars)+sep, coin_18p0,
        width=width, alpha=0.8, label='1e18.0 eV')
# Add a legend to the plot with the given labels
plt.legend(fontsize=10)
# Grab the current axes
#axes = plt.gca()
# Set the x tick positions to 0 - # of bars
#axes.set_xticks(np.arange(n_bars))
# Set the labels to be the strings proveded above in xlabels
#axes.set_xticklabels(xlabels)
# Set the fontsize of all axis labels
#for item in (axes.get_xticklabels() + axes.get_yticklabels()):
#    item.set_fontsize(12)
plt.xticks(ticks=np.arange(n_bars), labels=xlabels, fontsize=12)
plt.yticks(fontsize=12)
# Give the plot a title
plt.title("Surface Coincidence Fraction", fontsize=14)
# Save the plot as a PDF so it scales nicely in presentations
plt.savefig('station_surface_coincidence_fractions.pdf')



####################### Plot surface coincidence volumes ######################
xlabels = ["Surface", "15 m", "60 m", "100 m", "200 m"]
width = 0.4
sep = 0.2
n_bars = len(coin_vol_16p5)
# Set the size of the plot
plt.figure(figsize=(8, 5))
# Get the current axes
ax1 = plt.gca()
# Plot the first bar graph on the first axes
ax1.bar(np.arange(n_bars)-sep, coin_vol_16p5,
        width=width, color='C0', alpha=0.8,
        label='1e16.5 eV')
# Set up horizontal dotted grid lines (on the first axes)
plt.grid(axis='y', ls=':')
# Create secondary axes on top of the first that have the same x axis
ax2 = ax1.twinx()
# Plot the second bar graph on the second axes
plt.bar(np.arange(n_bars)+sep, coin_vol_18p0,
        width=width, color='C1', alpha=0.8,
        label='1e18.0 eV')
# Combine the legends of both axes by manually pulling the handles and labels
# from each individual axis and passing them to plt.legend()
ax1_handles, ax1_labels = ax1.get_legend_handles_labels()
ax2_handles, ax2_labels = ax2.get_legend_handles_labels()
plt.legend(ax1_handles+ax2_handles, ax1_labels+ax2_labels, fontsize=10)
# Set the label for the first y axis in its corresponding color
ax1.set_ylabel('1e16.5 eV', color='C0', fontsize=12)
# Set the color for the first y axis tick labels as well
for item in ax1.get_yticklabels():
    item.set_color('C0')
# Set the label for the second y axis in its corresponding color
ax2.set_ylabel('1e18.0 eV', color='C1', fontsize=12,
               rotation=270, verticalalignment='bottom')
# Set the color for the second y axis tick labels as well
for item in ax2.get_yticklabels():
    item.set_color('C1')
# Set the x tick positions to 0 - # of bars
ax1.set_xticks(np.arange(n_bars))
# Set the labels to be the strings proveded above in xlabels
ax1.set_xticklabels(xlabels)
# Set the fontsize of all axis labels
for item in (ax1.get_xticklabels() + ax1.get_yticklabels()):
    item.set_fontsize(12)
# Give the plot a title
plt.title("Surface-Coincident Events Effective Volume [km$^3$ sr]", fontsize=14)
# Save the plot as a PDF so it scales nicely in presentations
plt.savefig('coincidence_histogram.pdf')