import argparse
from numpy import load, ndarray
import os

parser = argparse.ArgumentParser()
parser.add_argument("-f", "--filepath", required=True, help="Specify the file path to the agent.", type=str)
args = parser.parse_args()

filepath = args.filepath
npdata = load(filepath)

evaluations = ndarray.tolist(npdata['results'])
# print(evaluations)
sorted_evals = []
for eval in evaluations:
    sorted_evals.append(sorted(eval))

# Now I have a sorted list.
# Now just pop the first and last elements of each eval
for eval in sorted_evals:
    eval.pop(0)
    eval.pop()

print()
# print(sorted_evals)

# Now that I have my sorted evaluations, I can calculate the mean episode reward for each eval
mean_eval_rewards = []
for eval in sorted_evals:
    mean_eval_rewards.append(sum(eval) / len(eval))
    
# Now I should have a list with the mean evaluation reward with the highest and lowest score tossed out
print(mean_eval_rewards)
print("num evals: " + str(len(mean_eval_rewards)))

# I'm dealing with a 2D array. Each element contains an array of ten data points
# The number of elements is going to vary for each training run
# The number of evaluation episodes will be constant, 10.
# I need to convert to a regular list first
# I could iterate over each element

agent_dirs = []
for d in os.listdir("agents/"):
    if "dqn_v2" in d:
        agent_dirs.append(d)
# Now I have a list of dirs with the evals.
# Iterate over the dirs, append the file path, load the evals, calculate the average score of the eval, then return a list with averages
eval_list = []
for d in agent_dirs:
    path = "agents/" + d + "/evaluations.npz"
    evals = ndarray.tolist(load(path)["results"])
    eval_list.append(evals)
# for i in eval_list:
#     print(i)
#     print()

def remove_outliers(evals: list) -> list:
    trimmed = []
    for eval in evals:
        eval.sort()
        eval.pop(0)
        eval.pop()
        trimmed.append(eval)
    return trimmed

avgs = [[]]
index = 0
for i in eval_list:
    avgs.append(i)
    for j in i:
        j.sort()
        j.pop()
        j.pop(0)
        avgs[index].append(sum(j) / len(j))
    index += 1
    
print(avgs)