import gym
from gym.envs.registration import register
import sys
import tty
import termios

# ================================================================================
# Class which is used to get "pressed keys"
class _Getch:
  def __call__(self):
    fd=sys.stdin.fileno()
    old_settings=termios.tcgetattr(fd)

    try:
      tty.setraw(sys.stdin.fileno())
      ch=sys.stdin.read(3)
    finally:
      termios.tcsetattr(fd,termios.TCSADRAIN,old_settings)
    return ch

# ================================================================================
# c inkey: instance of _Getch
inkey=_Getch()

# ================================================================================
# Configured key MACROS in OpenAI GYM

LEFT=0
DOWN=1
RIGHT=2
UP=3

# ================================================================================
# Key mapping from '\x1b[A' to UP

arrow_keys={
  '\x1b[A':UP,
  '\x1b[B':DOWN,
  '\x1b[C':RIGHT,
  '\x1b[D':LEFT}

# ================================================================================
# You configure FrozenLake environment with several additional configurations

register(
  id='FrozenLake-v3',
  entry_point='gym.envs.toy_text:FrozenLakeEnv',
  # (4,4) easy and simple grid
  kwargs={'map_name':'4x4','is_slippery':False})

# ================================================================================
# You create FrozenLake-v3 env

env = gym.make('FrozenLake-v3')

# ================================================================================
# Show the initial board visually

env.render()

# ================================================================================
while True:
  # Get pressed key
  key=inkey()
  if key not in arrow_keys.keys():
    print("Game aborted!")
    break
  
  # ================================================================================
  # Choose action from keyboard input
  action=arrow_keys[key]

  # ================================================================================
  # Put action into environment
  # And get state and reward
  # done: if game ends (like falling into hole, reaching to goal)
  state,reward,done,info=env.step(action)

  # ================================================================================
  # Show the board after action visually
  env.render()

  # ================================================================================
  print("State:",state,"Action:",action,"Reward:",reward,"Info:",info)

  # ================================================================================
  if done:
    print("Finished with reward",reward)
    break