import os

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
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
data_mean = pd.read_csv(os.path.join("data", "20201002_lifelong_final-acc-split-mnist_mean.csv"))
data_sem = pd.read_csv(os.path.join("data", "20201002_lifelong_final-acc-split-mnist_sem.csv"))

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

# Plot the bar chart
data_mean.drop(columns='Joint Training').plot.bar(ax=ax, rot=0, yerr=data_sem, error_kw={"capsize": 1, "capthick": 0.4, "elinewidth": 0.4})
ax.axhline(data_mean['Joint Training'].max(), color='black', linestyle='--')

# Shade area behind bars depicting average values
x_anchor = (ax.get_xticks()[-1] + ax.get_xticks()[-2]) / 2
ax.axvline(x_anchor, color='black', linewidth=0.1)
ax.axvspan(x_anchor, 99, alpha=0.05, color='black', linewidth=0.0)

# X/Y Labels and ticks
ax.set_xticks(np.arange(0, 5 + 1, 1))
ax.set_xticklabels(list(range(1, 6)) + ["All Tasks"])
ax.set_ylim(0.7, 1.0)
ax.set_xlabel("Task")
ax.set_ylabel("Test Accuracy")

# Add the legend
ax.legend(ncol=2, bbox_to_anchor=(0.5, 1.0))

# Save the figure
fig.tight_layout(pad=0.0)
fig.savefig("lifelong-mlp_bar_final-acc_all.pdf")
