from dm_control import suite
from dm_control.suite.wrappers import pixels
from dm_env.specs import Array, BoundedArray

import numpy as np
import os
import copy
from collections import namedtuple, OrderedDict
from cfm.env.utils import namedarraytuple, Env, EnvStep, EnvSpaces
from cfm.env.spaces.box import Box
from cfm.env.spaces.composite import Composite

State = None

def convert_dm_control_to_rlpyt_space(dm_control_space):
    """Recursively convert dm_control_space into gym space.

    Note: Need to check the following cases of the input type, in the following
    order:
       (1) BoundedArray
       (2) Array
       (3) OrderedDict.

    - Generally, dm_control observation_specs are OrderedDict with other spaces
      (e.g. Array) nested in it.
    - Generally, dm_control action_specs are of type `BoundedArray`.

    To handle dm_control observation_specs as inputs, we check the following
    input types in order to enable recursive calling on each nested item.
    """
    if isinstance(dm_control_space, BoundedArray):
        rlpyt_box = Box(
            low=dm_control_space.minimum,
            high=dm_control_space.maximum,
            shape=None,
            dtype=dm_control_space.dtype)
        assert rlpyt_box.shape == dm_control_space.shape, (
            (rlpyt_box.shape, dm_control_space.shape))
        return rlpyt_box
    elif isinstance(dm_control_space, Array):
        if isinstance(dm_control_space, BoundedArray):
            raise ValueError("The order of the if-statements matters.")
        return Box(
            low=-float("inf"),
            high=float("inf"),
            shape=dm_control_space.shape,
            dtype=dm_control_space.dtype)
    elif isinstance(dm_control_space, OrderedDict):
        global State
        if State is None:
            State = namedtuple('State', list(dm_control_space.keys()))
        return Composite([convert_dm_control_to_rlpyt_space(value)
                          for value in dm_control_space.values()], State)
    else:
        raise ValueError(dm_control_space)

EnvInfo = None
Observation = None

def init_namedtuples(info_keys=None, state_keys=None):
    global EnvInfo, Observation, State

    if info_keys is None:
        info_keys = ['traj_done']

    if state_keys is None:
        state_keys = ['pixels']

    EnvInfo = namedtuple('EnvInfo', info_keys)
    Observation = namedarraytuple('Observation', state_keys)
    State = namedtuple('State', state_keys)

class DMControlEnv(Env):

    def __init__(self,
                 domain,
                 task,
                 frame_skip=1,
                 normalize=False,
                 pixel_wrapper_kwargs=None,
                 task_kwargs={},
                 environment_kwargs={},
                 max_path_length=1200,
                 ):
        self._max_path_length = max_path_length

        env = suite.load(domain_name=domain,
                         task_name=task,
                         task_kwargs=task_kwargs,
                         environment_kwargs=environment_kwargs)
        if normalize:
            np.testing.assert_equal(env.action_spec().minimum, -1)
            np.testing.assert_equal(env.action_spec().maximum, 1)
        if pixel_wrapper_kwargs is not None:
            env = pixels.Wrapper(env, **pixel_wrapper_kwargs)
        self._env = env

        self._observation_keys = tuple(env.observation_spec().keys())
        observation_space = convert_dm_control_to_rlpyt_space(
            env.observation_spec())
        self._observation_space = observation_space

        action_space = convert_dm_control_to_rlpyt_space(env.action_spec())
        if len(action_space.shape) > 1:
            raise NotImplementedError(
                "Shape of the action space ({}) is not flat, make sure to"
                " check the implemenation.".format(action_space))
        self._action_space = action_space

        self._step_count = 0

    def reset(self):
        self._step_count = 0
        time_step = self._env.reset()
        observation = self._filter_observation(time_step.observation)

        global Observation
        if Observation is None:
            Observation = namedarraytuple("Observation", list(observation.keys()))
        observation = Observation(**{k: v for k, v in observation.items()
                                     if k in self._observation_keys})
        return observation

    def step(self, action):
        time_step = self._env.step(action)
        reward = time_step.reward
        terminal = time_step.last()
        info = time_step.info
        info.update({
            key: value
            for key, value in time_step.observation.items()
            if key not in self._observation_keys
        })
        observation = self._filter_observation(time_step.observation)

        self._step_count += 1
        info['traj_done'] = self._step_count >= self._max_path_length

        global EnvInfo
        if EnvInfo is None:
            EnvInfo = namedtuple("EnvInfo", list(info.keys()))
        info = EnvInfo(**{k: v for k, v in info.items() if k in EnvInfo._fields})

        global Observation
        if Observation is None:
            Observation = namedarraytuple("Observation", list(observation.keys()))
        observation = Observation(**{k: v.copy() for k, v in observation.items()
                                     if k in self._observation_keys})

        return EnvStep(observation, reward, terminal, info)

    def render(self, *args, mode='rgb_array', width=256, height=256,
               cameria_id=0, **kwargs):
        if mode == 'human':
            raise NotImplementedError(
                "TODO(Alacarter): Figure out how to not continuously launch"
                " viewers if one is already open."
                " See: https://github.com/deepmind/dm_control/issues/39.")
        elif mode == 'rgb_array':
            return self._env.physics.render(width=width, height=height,
                                            camera_id=cameria_id, **kwargs)
        raise NotImplementedError(mode)

    def get_obs(self):
        obs = self._env.task.get_observation(self._env.physics)
        obs['pixels'] = self._env.physics.render(**self._env._render_kwargs)
        obs = self._filter_observation(obs)
        obs = Observation(**{k: v for k, v in obs.items()
                             if k in self._observation_keys})
        return obs

    def get_state(self, ignore_step=True):
        if ignore_step:
            return self._env.physics.get_state()
        return self._env.physics.get_state(), self._step_count

    def set_state(self, state, ignore_step=True):
        if ignore_step:
            self._env.physics.set_state(state)
            self._env.step(np.zeros(self.action_space.shape))
        else:
            self._env.physics.set_state(state[0])
            self._env.step(np.zeros(self.action_space.shape))
            self._step_count = state[1]

    def get_geoms(self):
        return self._env.task.get_geoms(self._env.physics)

    @property
    def spaces(self):
        return EnvSpaces(
            observation=self._observation_space,
            action=self._action_space,
        )

    ###########################################################################
    # Helpers

    def _filter_observation(self, observation):
        observation = type(observation)([
            (name, value)
            for name, value in observation.items()
            if name in self._observation_keys
        ])
        return observation

    ###########################################################################
    # Properties
