#!/usr/bin/env python3

import time
import argparse
import numpy as np
import gym
import gym_minigrid
from gym_minigrid.wrappers import *
from gym_minigrid.window import Window
from envs.env import *
from dqn.library import *
import itertools

def redraw(img):
    if not args.agent_view:
        img = env.render('rgb_array', tile_size=args.tile_size, highlight=False)

    window.show_img(img)

def reset():
    if args.seed != -1:
        env.seed(args.seed)

    obs = env.reset()

    if hasattr(env, 'mission'):
        print('Mission: %s' % env.mission)
        window.set_caption(env.mission)

    redraw(obs)

def step(action):
    obs, reward, done, info = env.step(action)
    print('step=%s, reward=%.2f' % (env.step_count, reward))

    if done:
        print('done! ', to_hash(obs))
        reset()
    else:
        redraw(obs)

def key_handler(event):
    print('pressed', event.key)

    if event.key == 'escape':
        window.close()
        return

    if event.key == 'backspace':
        reset()
        return

    if event.key == 'left':
        step(env.actions.left)
        return
    if event.key == 'right':
        step(env.actions.right)
        return
    if event.key == 'up':
        step(env.actions.forward)
        return

    # Spacebar
    if event.key == ' ':
        step(env.actions.toggle)
        return
    if event.key == 'pageup':
        step(env.actions.pickup)
        return
    if event.key == 'pagedown':
        step(env.actions.drop)
        return

    if event.key == 'enter':
        step(env.actions.done)
        return

parser = argparse.ArgumentParser()
parser.add_argument(
    '--obj_color',
    default=None,
    help="Object color"
)
parser.add_argument(
    '--obj_type',
    default=None,
    help="Object type"
)
parser.add_argument(
    '--room_num',
    type=int,
    default=None,
    help="Room number"
)
parser.add_argument(
    '--size',
    type=int,
    default=7,
    help="Room number"
)
parser.add_argument(
    '--dist_type',
    default=None,
    help="Distractor type"
)
parser.add_argument(
    '--dist_color',
    default=None,
    help="Distractor color"
)
parser.add_argument(
    '--num_dists',
    type=int,
    default=4,
    help="Number of distractors"
)
parser.add_argument(
    "--seed",
    type=int,
    help="random seed to generate the environment with",
    default=-1
)
parser.add_argument(
    "--tile_size",
    type=int,
    help="size at which to render tiles",
    default=32
)
parser.add_argument(
    '--agent_view',
    default=False,
    help="draw the agent sees (partially observable view)",
    action='store_true'
)

args = parser.parse_args()

env = make_env()
if not (args.obj_color or args.obj_type or args.room_num):
    goals = env.all_goals
elif not (args.obj_color or args.obj_type):
    goals = list(itertools.product(OBJ_COLORS,OBJ_TYPES,[env.rooms[args.room_num-1]]))
elif not (args.obj_color or args.room_num):
    goals = list(itertools.product(OBJ_COLORS,[args.obj_type],env.rooms))
elif not (args.obj_type or args.room_num):
    goals = list(itertools.product([args.obj_color],OBJ_TYPES,env.rooms))
elif not args.room_num:
    goals = list(itertools.product([args.obj_color],[args.obj_type],env.rooms))
elif not args.obj_type:
    goals = list(itertools.product([args.obj_color],OBJ_TYPES,[env.rooms[args.room_num-1]]))
elif not args.obj_color:
    goals = list(itertools.product(OBJ_COLORS,[args.obj_type],[env.rooms[args.room_num-1]]))

print('goals: ',len(goals))
# goals = list(itertools.product(['grey'],['red','green']))
env = make_env(size=args.size, goals=goals,
               dist_type=args.dist_type, dist_color=args.dist_color, num_dists=args.num_dists)

window = Window('gym_minigrid')
window.reg_key_handler(key_handler)

reset()

# Blocking event loop
window.show(block=True)
