from numpy import load
import matplotlib.pyplot as plt 
# import matplotlib.axes

filepath = "dqn_v2-5/evaluations.npz"

data = load(filepath)
lst = data.files # data.files lists the keys that are available for data

# print('ep_lengths: \n', data['ep_lengths'])

# results and ep_lengths are 2d arrays, because each evaluation is 5 episodes long.
# I want to plot the average of each evaluation.

# for each item in results, loop through the array and save the average
avg_ep_result_arr = []
for eval in data['results']:
    result_sum = 0
    
    for result in eval:
        result_sum = result_sum + result
        
    avg_ep_result = result_sum / len(eval)
    avg_ep_result_arr.append(avg_ep_result)
    
avg_ep_len_arr = []
for eval in data['ep_lengths']:
    max_len = 0
    y_limit = 0
    ep_len_sum = 0
    
    for ep_length in eval:
        ep_len_sum = ep_len_sum + ep_length
        if ep_length > max_len:
            max_len = ep_length
        if ep_length > y_limit and y_limit < max_len:
            y_limit = ep_length
        
    avg_ep_len = ep_len_sum / len(eval)
    avg_ep_len_arr.append(avg_ep_len)
    y_limit = y_limit * 1.9
    
    
x = plt.plot(data['timesteps'], avg_ep_result_arr)
# plt.bar(data['timesteps'], avg_ep_len_arr, width=10000)
y = plt.plot(data['timesteps'], avg_ep_len_arr)
plt.ylim(top=y_limit)
# plt.ylabel("Avg ep score")
# lineObjects = plt.plot(x, y)
plt.legend(["avg ep result", "avg ep length"])
plt.title("result and length over steps\nfilepath: " + filepath)
plt.show()