from __future__ import print_function
import math
import sys
import gym
import gym.spaces
import numpy as np
from gym import core, spaces
from gym.utils import seeding
from numpy import sin, cos, pi
import time

env = gym.make('Pendulum-v0')
ROLLOUT_TIME = 1000
SKIP_CONTROL = 0
human_agent_action = 0
human_wants_restart = False
human_sets_pause = False

def key_press(key, mod):
    global human_agent_action, human_wants_restart, human_sets_pause
    if key == 0xff0d:
        human_wants_restart = True
    if key == 32:
        human_sets_pause = not human_sets_pause
    a = int(key - ord('0'))
    if a == 0:
        a = -1
    if a == 1:
        a = 1
    human_agent_action = a

def key_release(key, mod):
    global human_agent_action
    a = int(key - ord('0'))
    if a == 0:
        a = -1
    if a == 1:
        a = 1
    if human_agent_action == a:
        human_agent_action = 0
env.reset()
env.render()
env.unwrapped.viewer.window.on_key_press = key_press
env.unwrapped.viewer.window.on_key_release = key_release

def rollout(env):
    global human_agent_action, human_wants_restart, human_sets_pause
    human_wants_restart = False
    obser = env.reset()
    skip = 0
    for t in range(ROLLOUT_TIME):
        if not skip:
            a = human_agent_action
            skip = SKIP_CONTROL
        else:
            skip -= 1
        print(a)
        obser, r, done, info = env.step([a])
        env.render()
        if done:
            break
        if human_wants_restart:
            break
while 1:
    rollout(env)