# Copyright 2017 The dm_control Authors.
#
# 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.
# ============================================================================

"""Base class for tasks in the Control Suite."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from dm_control import mujoco
from dm_control.rl import control

import numpy as np


class Task(control.Task):
    """Base class for tasks in the Control Suite.

    Actions are mapped directly to the states of MuJoCo actuators: each element of
    the action array is used to set the control input for a single actuator. The
    ordering of the actuators is the same as in the corresponding MJCF XML file.

    Attributes:
      random: A `numpy.random.RandomState` instance. This should be used to
        generate all random variables associated with the task, such as random
        starting states, observation noise* etc.

    *If sensor noise is enabled in the MuJoCo model then this will be generated
    using MuJoCo's internal RNG, which has its own independent state.
    """

    def __init__(self, random=None):
        """Initializes a new continuous control task.

        Args:
          random: Optional, either a `numpy.random.RandomState` instance, an integer
            seed for creating a new `RandomState`, or None to select a seed
            automatically (default).
        """
        if not isinstance(random, np.random.RandomState):
            random = np.random.RandomState(random)
        self._random = random
        self._visualize_reward = False

    @property
    def random(self):
        """Task-specific `numpy.random.RandomState` instance."""
        return self._random

    def action_spec(self, physics):
        """Returns a `BoundedArraySpec` matching the `physics` actuators."""
        return mujoco.action_spec(physics)

    def initialize_episode(self, physics):
        """Resets geom colors to their defaults after starting a new episode.

        Subclasses of `base.Task` must delegate to this method after performing
        their own initialization.

        Args:
          physics: An instance of `mujoco.Physics`.
        """
        self.after_step(physics)

    def before_step(self, action, physics):
        """Sets the control signal for the actuators to values in `action`."""
        # Support legacy internal code.
        action = getattr(action, "continuous_actions", action)
        physics.set_control(action)

    def after_step(self, physics):
        """Modifies colors according to the reward."""
        if self._visualize_reward:
            reward = np.clip(self.get_reward(physics), 0.0, 1.0)
            _set_reward_colors(physics, reward)

    @property
    def visualize_reward(self):
        return self._visualize_reward

    @visualize_reward.setter
    def visualize_reward(self, value):
        if not isinstance(value, bool):
            raise ValueError("Expected a boolean, got {}.".format(type(value)))
        self._visualize_reward = value


_MATERIALS = ["self", "effector", "target"]
_DEFAULT = [name + "_default" for name in _MATERIALS]
_HIGHLIGHT = [name + "_highlight" for name in _MATERIALS]


def _set_reward_colors(physics, reward):
    """Sets the highlight, effector and target colors according to the reward."""
    assert 0.0 <= reward <= 1.0
    colors = physics.named.model.mat_rgba
    default = colors[_DEFAULT]
    highlight = colors[_HIGHLIGHT]
    blend_coef = reward ** 4  # Better color distinction near high rewards.
    colors[_MATERIALS] = blend_coef * highlight + (1.0 - blend_coef) * default
