import numpy as np
import matplotlib.pyplot as plt
# package to load csv files
from numpy import genfromtxt
# used in customized legends
import matplotlib.patches as mpatches
from matplotlib.lines import Line2D

# load csv files with numpy
reward = genfromtxt('reward.csv', delimiter=',')
reward_lst = reward[1:,2]
# reward_lst: some one-dimensional array
reward_sd = genfromtxt('reward_sd.csv', delimiter=',')
reward_sd_lst = reward_sd[1:,2]
# reward_sd_lst: some one-dimensional array

plt.figure()
x = [500*i for i in range(60)]
y_x = [200 for _ in x]

# plotting the line and fill-in area
plt.plot(x, reward_lst)
plt.plot(x, y_x, '--', c='grey')
plt.fill_between(x, reward_lst-reward_sd_lst, reward_lst+reward_sd_lst, alpha=0.2)

# labels and titles
plt.title('Testing epoch reward')
plt.ylabel('accumulative reward')
plt.xlabel('epoch number')

# customize the legend bar
blue_patch = mpatches.Patch(color=u'#1f77b4', label='cumulative reward ± 1 sd', alpha=0.2)
line1 = Line2D([0], [0], color=u'#1f77b4', label='mean cumulative reward')
line2 = Line2D([0], [0], color='grey', linestyle='--', label='reference level (200)')
plt.legend(handles=[blue_patch, line1, line2], loc='lower right')

# save figure
plt.savefig('testing_reward.png', dpi=600)
plt.show()