import gym
import pandas as pd
import seaborn as sns
from tqdm import tqdm
from rlagents import QLAgent, benchmark, gymwrapper

env_list = [gymwrapper(gym.make('CartPole-v0'))]
agent_list = [QLAgent()]

paths = benchmark(agent_list, env_list,
                  n_episodes=500,
                  max_steps_per_episode=200,
                  n_trials=200)

paths = ['DataDir/'+i for i in os.listdir('DataDir/')]
data = []
for p in tqdm(paths, leave=False):
	with open(p, 'r') as fl:
		da = ujson.load(fl)

		x = [(None, 'Monte' in d[1], None, int(d[3]), float(d[4]), int(d[5]))
			 for d in da]
		data.extend(x)

data = pd.DataFrame(data, columns=['runid', 'agent', 'world', 'ep', 'rew', 'trial'])
data.agent = 'Q Learning'
sns.tsplot(data, time='ep', value='rew', unit='trial', condition='agent', ci=95)