import os

import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from scipy.stats import pearsonr

# Global configuration of matplotlib to use custom style sheet
matplotlib.use('pgf')
matplotlib.style.use("style.mplstyle")

# Import the data from the csv file
# path = os.path.join("data", "20200124_probs-fisher_mnist.csv")
path = os.path.join("data", "20200206_energy_probs-fisher.csv")
data = pd.read_csv(path)
data["Important"] = data['Transmission Probability'] > 0.9

# Prepare the figure
fig, ax = plt.subplots(figsize=[2.6, 1.5])

# Compute the Pearson correlation
corr_coeff, p_value = pearsonr(data["Transmission Probability"], data["Fisher Information"])

# Plot the histogram
ax = sns.violinplot(x="Important", y="Fisher Information", data=data, bw=1.0, cut=0, scale='count', ax=ax, inner=None, linewidth=0.0)

# Add annotation with correlation coefficient
# ax.annotate("$r={:.4f}$".format(corr_coeff), (0.2, 0.9), xycoords='axes fraction', ha='center', va='center', bbox=dict(facecolor='white', edgecolor='none', boxstyle='square'))

# X/Y labels and ticks
ax.set_xticklabels(['Low\nRelease Probability', 'High\nRelease Probability'])
ax.set_ylim(-0.001, 0.03)
ax.set_xlabel(None)
ax.set_ylabel("Fisher Information")

# Save the figure
fig.tight_layout(pad=0.0)
fig.savefig("energy-mlp_violin_fisher.pdf")
