import cv2
from gym.spaces.box import Box
import numpy as np
import gym
from gym import spaces
import logging
import universe
from universe import vectorized
from universe.wrappers import BlockingReset, GymCoreAction, EpisodeID, Unvectorize, Vectorize, Vision, Logger
from universe import spaces as vnc_spaces
from universe.spaces.vnc_event import keycode
import time
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
universe.configure_logging()

import os
import sys

this_dir = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, os.path.join(this_dir, '..'))
from envs import AtariRescale42x42, CropScreen, DiagnosticsInfo, DiscreteToFixedKeysVNCActions, FlashRescale


def create_env(env_id, client_id, remotes, **kwargs):
  spec = gym.spec(env_id)

  if spec.tags.get('flashgames', False):
    return create_flash_env(env_id, client_id, remotes, **kwargs)
  elif spec.tags.get('atari', False) and spec.tags.get('vnc', False):
    return create_vncatari_env(env_id, client_id, remotes, **kwargs)
  else:
    # Assume atari.
    assert "." not in env_id  # universe environments have dots in names.
    return create_atari_env(env_id)


def create_flash_env(env_id, client_id, remotes, **_):
  env = gym.make(env_id)
  env = Vision(env)
  env = Logger(env)
  env = BlockingReset(env)

  reg = universe.runtime_spec('flashgames').server_registry
  height = reg[env_id]["height"]
  width = reg[env_id]["width"]
  env = CropScreen(env, height, width, 84, 18)
  env = FlashRescale(env)

  keys = ['left', 'right', 'up', 'down', 'x']
  if env_id == 'flashgames.NeonRace-v0':
    # Better key space for this game.
    keys = ['left', 'right', 'up', 'left up', 'right up', 'down', 'up x']
  logger.info('create_flash_env(%s): keys=%s', env_id, keys)

  env = DiscreteToFixedKeysVNCActions(env, keys)
  env = EpisodeID(env)
  env = DiagnosticsInfo(env)
  env = Unvectorize(env)
  env.configure(
      fps=5.0,
      remotes=remotes,
      start_timeout=15 * 60,
      client_id=client_id,
      vnc_driver='go',
      vnc_kwargs={
          'encoding': 'tight',
          'compress_level': 0,
          'fine_quality_level': 50,
          'subsample_level': 3
      })
  return env


def create_vncatari_env(env_id, client_id, remotes, **_):
  env = gym.make(env_id)
  env = Vision(env)
  env = Logger(env)
  env = BlockingReset(env)
  env = GymCoreAction(env)
  env = AtariRescale42x42(env)
  env = EpisodeID(env)
  env = DiagnosticsInfo(env)
  env = Unvectorize(env)

  logger.info('Connecting to remotes: %s', remotes)
  fps = env.metadata['video.frames_per_second']
  env.configure(
      remotes=remotes, start_timeout=15 * 60, fps=fps, client_id=client_id)
  return env


def create_atari_env(env_id):
  env = gym.make(env_id)
  #    env = Vectorize(env)
  #    env = AtariRescale42x42(env)
  #    env = DiagnosticsInfo(env)
  #    env = Unvectorize(env)
  return env
