import os

import matplotlib
import matplotlib.pyplot as plt
import pandas as pd

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

# Import the data from the respective csv files
keys = ['lesion_prob', 'lesion_rand']
data_mean = pd.read_csv(os.path.join("data", "20201007_energy_lesion_mean.csv"), index_col='Fraction')[keys]
data_sem = pd.read_csv(os.path.join("data", "20201007_energy_lesion_sem.csv"), index_col='Fraction')[keys]

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

# Plot the line plots
data_mean.plot.line(ax=ax)

# Plot the confidence intervals for all columns
for i, key in enumerate(data_mean.columns):
    lowerbound = data_mean[key] - data_sem[key]
    upperbound = data_mean[key] + data_sem[key]
    ax.fill_between(x=data_mean.index, y1=lowerbound, y2=upperbound, alpha=0.3)

# Add the legend
ax.legend(ncol=1, loc="lower left", frameon=True, labels=["Lowest Probability First", "Random Order"])

# X/Y Labels and ticks
ax.set_xlim(0.0, 0.9)
ax.set_ylim(0.0, 1.0)
ax.set_xlabel("Fraction of Lesioned Weights")
ax.set_ylabel('Accuracy')

# Save the figure
plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
fig.tight_layout(pad=0.0)
fig.savefig("energy-mlp_line_importance-lesion.pdf")
