import numpy as np
import gym

from IPython.display import display

import qcircuit

env = gym.make('qcircuit-v1')

env.reset()
display(env.render())

done = False
while(not done):
    obs, _, done, info = env.step(env.action_space.sample())
    display(info['circuit_img'])
       
env.close()

from stable_baselines.common.policies import MlpPolicy
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2

env = DummyVecEnv([lambda: env])
modelPPO2 = PPO2(MlpPolicy, env, verbose=1)
modelPPO2.learn(total_timesteps=10000)

obs = env.reset()
display(env.render())

for _ in range(10):
    action, _states = modelPPO2.predict(obs)
    obs, _, done, info = env.step(action)
    display(info[0]['circuit_img'])
    
env.close()

from stable_baselines import A2C

modelA2C = A2C(MlpPolicy, env, verbose=1)
modelA2C.learn(total_timesteps=10000)

obs = env.reset()
display(env.render())

for _ in range(10):
    action, _states = modelA2C.predict(obs)
    obs, _, done, info = env.step(action)
    display(info[0]['circuit_img'])
    
env.close()

import evaluation
n_episodes = 1000

PPO2_perf, _ = evaluation.evaluate_model(modelPPO2, env, num_steps=n_episodes)
A2C_perf, _ = evaluation.evaluate_model(modelA2C, env, num_steps=n_episodes)

env = gym.make('qcircuit-v1')
rand_perf, _ = evaluation.evaluate_random(env, num_steps=n_episodes)

print('Mean performance of random agent (out of {0} episodes): {1}'.format(n_episodes,rand_perf))
print('Mean performance of PPO2 agent (out of {0} episodes): {1}'.format(n_episodes,PPO2_perf))
print('Mean performance of A2C agent (out of {0} episodes): {1}'.format(n_episodes,A2C_perf))