#
# Copyright (c) 2017 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import sys
from logger import *
import gym
import numpy as np
import time
try:
    import roboschool
    from OpenGL import GL
except ImportError:
    from logger import failed_imports
    failed_imports.append("RoboSchool")

try:
    from gym_extensions.continuous import mujoco
except:
    from logger import failed_imports
    failed_imports.append("GymExtensions")

try:
    import pybullet_envs
except ImportError:
    from logger import failed_imports
    failed_imports.append("PyBullet")

from gym import wrappers
from utils import force_list, RunPhase
from environments.environment_wrapper import EnvironmentWrapper


class GymEnvironmentWrapper(EnvironmentWrapper):
    def __init__(self, tuning_parameters):
        EnvironmentWrapper.__init__(self, tuning_parameters)

        # env parameters
        if ':' in self.env_id:
            self.env = gym.envs.registration.load(self.env_id)()
        else:
            self.env = gym.make(self.env_id)

        if self.seed is not None:
            self.env.seed(self.seed)

        # self.env_spec = gym.spec(self.env_id)
        self.env.frameskip = self.frame_skip
        self.discrete_controls = type(self.env.action_space) != gym.spaces.box.Box

        self.state = self.reset(True)['state']

        # render
        if self.is_rendered:
            image = self.get_rendered_image()
            scale = 1
            if self.human_control:
                scale = 2
            self.renderer.create_screen(image.shape[1]*scale, image.shape[0]*scale)

        # TODO: collect and store this as observation space instead
        self.is_state_type_image = len(self.state['observation'].shape) > 1
        if self.is_state_type_image:
            self.width = self.state['observation'].shape[1]
            self.height = self.state['observation'].shape[0]
        else:
            self.width = self.state['observation'].shape[0]

        # action space
        self.actions_description = {}
        if hasattr(self.env.unwrapped, 'get_action_meanings'):
            self.actions_description = self.env.unwrapped.get_action_meanings()
        if self.discrete_controls:
            self.action_space_size = self.env.action_space.n
            self.action_space_abs_range = 0
        else:
            self.action_space_size = self.env.action_space.shape[0]
            self.action_space_high = self.env.action_space.high
            self.action_space_low = self.env.action_space.low
            self.action_space_abs_range = np.maximum(np.abs(self.action_space_low), np.abs(self.action_space_high))
        self.actions = {i: i for i in range(self.action_space_size)}
        self.key_to_action = {}
        if hasattr(self.env.unwrapped, 'get_keys_to_action'):
            self.key_to_action = self.env.unwrapped.get_keys_to_action()

        # measurements
        if self.env.spec is not None:
            self.timestep_limit = self.env.spec.timestep_limit
        else:
            self.timestep_limit = None
        self.measurements_size = len(self.step(0)['info'].keys())

    def _wrap_state(self, state):
        if isinstance(self.env.observation_space, gym.spaces.Dict):
            return state
        else:
            return {'observation': state}

    def _update_state(self):
        if hasattr(self.env, 'env') and hasattr(self.env.env, 'ale'):
            if self.phase == RunPhase.TRAIN and hasattr(self, 'current_ale_lives'):
                # signal termination for life loss
                if self.current_ale_lives != self.env.env.ale.lives():
                    self.done = True
            self.current_ale_lives = self.env.env.ale.lives()

    def _take_action(self, action_idx):
        if action_idx is None:
            action_idx = self.last_action_idx

        if self.discrete_controls:
            action = self.actions[action_idx]
        else:
            action = action_idx

        # pendulum-v0 for example expects a list
        if not self.discrete_controls:
            # catching cases where the action for continuous control is a number instead of a list the
            # size of the action space
            if type(action_idx) == int and action_idx == 0:
                # deal with the "reset" action 0
                action = [0] * self.env.action_space.shape[0]
            action = np.array(force_list(action))
            # removing redundant dimensions such that the action size will match the expected action size from gym
            if action.shape != self.env.action_space.shape:
                action = np.squeeze(action)
            action = np.clip(action, self.action_space_low, self.action_space_high)

        state, self.reward, self.done, self.info = self.env.step(action)
        self.state = self._wrap_state(state)

    def _preprocess_state(self, state):
        # TODO: move this into wrapper
        if any(env in self.env_id for env in ["Breakout", "Pong"]):
            # crop image
            state['observation'] = state['observation'][34:195, :, :]
        return state

    def _restart_environment_episode(self, force_environment_reset=False):
        # prevent reset of environment if there are ale lives left
        if (hasattr(self.env, 'env') and hasattr(self.env.env, 'ale') and self.env.env.ale.lives() > 0) \
                and not force_environment_reset and not self.env._past_limit():
            return self.state

        if self.seed:
            self.env.seed(self.seed)

        self.state = self._wrap_state(self.env.reset())
        while self.state is None:
            self.step(0)

        return self.state

    def get_rendered_image(self):
        return self.env.render(mode='rgb_array')

@@ -2099,24 +2099,630 @@
 %5B0%5D*scale)%0A%0A
+        if isinstance(self.env.observation_space, gym.spaces.Dict):%0A            if 'observation' not in self.env.observation_space:%0A                raise ValueError((%0A                    'The gym environment provided %7Benv_id%7D does not contain '%0A                    '%22observation%22 in its observation space. For now this is '%0A                    'required. The environment does include the following '%0A                    'keys in its observation space: %7Bkeys%7D'%0A                ).format(%0A                    env_id=self.env_id,%0A                    keys=self.env.observation_space.keys(),%0A                ))%0A%0A
         # TO

make sure that for now observation spaces all include an observation key