import gym
import gym_sokoban
from gym_sokoban.envs.room_utils import ACTION_LOOKUP
import time
from PIL import Image
import numpy as np
import argparse
import os

parser = argparse.ArgumentParser(description='Run environment with random selected actions.')
parser.add_argument('--rounds', '-r', metavar='rounds', type=int,
                    help='number of rounds to play (default: 1)', default=1)
parser.add_argument('--steps', '-s', metavar='steps', type=int,
                    help='maximum number of steps to be played each round (default: 300)', default=300)
parser.add_argument('--env', '-e', metavar='env',
                    help='Environment to load (default: Sokoban-v0)', default='Sokoban-v0')
parser.add_argument('--save', action='store_true',
                    help='Save images of single steps')

args = parser.parse_args()
env_name = args.env
n_rounds = args.rounds
n_steps = args.steps
save_images = args.save

# Creating target directory if images are to be stored
if save_images and not os.path.exists('images'):
    try:
        os.makedirs('images')
    except OSError:
        print('Error: Creating images target directory. ')

ts = time.time()
env = gym.make(env_name)


def print_avilable_actions():
    """
    Prints all available actions nicely formatted..
    :return:
    """
    available_actions_list = []
    for i in range(len(ACTION_LOOKUP)):
        available_actions_list.append(
            'Key: {} - Action: {}'.format(i, ACTION_LOOKUP[i])
        )
    display_actions = '\n'.join(available_actions_list)
    print()
    print('Action out of Range!')
    print('Available Actions:\n{}'.format(display_actions))
    print()


for i_episode in range(n_rounds):
    print('Starting new game!')
    observation = env.reset()

    for t in range(n_steps):
        env.render()

        action = input('Select action: ')
        try:
            action = int(action)

            if not action in range(len(ACTION_LOOKUP)):
                raise ValueError

        except ValueError:
            print_avilable_actions()
            continue

        observation, reward, done, info = env.step(action)
        print(ACTION_LOOKUP[action], reward, done, info)

        if save_images:
            img = Image.fromarray(np.array(observation), 'RGB')
            img.save(os.path.join('images', 'observation_{}_{}.png'.format(i_episode, t)))

        if done:
            print("Episode finished after {} timesteps".format(t+1))
            env.render()
            break


time.sleep(10)

@@ -810,16 +810,126 @@
  steps')
+%0Aparser.add_argument('--gifs', action='store_true',%0A                    help='Generate Gif files from images')
 %0A%0Aargs =
@@ -1036,16 +1036,55 @@
 rgs.save
+ or args.gifs%0Agenerate_gifs = args.gifs
 %0A%0A# Crea
@@ -2655,17 +2655,552 @@
  break%0A%0A
-%0A
+env.close()%0A%0Aif generate_gifs:%0A    print('')%0A    import imageio%0A%0A    for i_episode in range(n_rounds):%0A%0A        with imageio.get_writer(os.path.join('images', 'round_%7B%7D.gif'.format(i_episode)), mode='I') as writer:%0A%0A            for t in range(n_steps):%0A                try:%0A%0A                    filename = os.path.join('images', 'observation_%7B%7D_%7B%7D.png'.format(i_episode, t))%0A                    image = imageio.imread(filename)%0A                    writer.append_data(image)%0A%0A                except:%0A                    break%0A%0Aelse:%0A    
 time.sle
@@ -3205,9 +3205,8 @@
 leep(10)
-%0A

Add gif generation support for Human playing