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")
colors = ['#91056A', '#6F6F6F']

# Load the comparison data as DataFrames
data_mean = pd.read_csv(os.path.join("data", "20210104_lifelong_time-task1_acc_mean.csv"))
data_sem = pd.read_csv(os.path.join("data", "20210104_lifelong_time-task1_acc_sem.csv"))

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

# Plot the data as a line plot
data_mean[['Presynaptic Consolidation', 'No Consolidation']].plot(ax=ax, color=colors)
ax.axhline(data_mean['Joint Training'].max(), color='black', linestyle='--')


# Add confidence interval
for i, key in enumerate(['Presynaptic Consolidation', 'No Consolidation']):
    x = np.arange(0, len(data_mean[key]), 1)
    lowerbound = data_mean[key] - data_sem[key]
    upperbound = data_mean[key] + data_sem[key]
    ax.fill_between(x=x, y1=lowerbound, y2=upperbound, alpha=0.3, color=colors[i], lw=0.0)

# X/Y Labels and ticks
ax.set_xticks(np.arange(0, 5, 1))
ax.set_xticklabels(range(1, 6))
ax.set_ylim(0.5, 1.0)
ax.set_xlim(0.0, 4.0)
ax.set_xlabel("Tasks Trained On")
ax.set_ylabel("Test Accuracy Task 1")

# Add the legend
ax.legend(ncol=1, loc="lower center", frameon=True)

# Save the figure
fig.tight_layout(pad=0.0)
fig.savefig("lifelong-mlp_line_time-task1-acc.pdf")
