# Copyright 2018-2019 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.
# ============================================================================
"""Environment's execution runtime."""

import collections
import copy
import enum

from dm_control.mujoco.wrapper import mjbindings
from dm_control.viewer import util
import mujoco
import numpy as np

mjlib = mjbindings.mjlib


# Pause interval between simulation steps.
_SIMULATION_STEP_INTERVAL = 0.001

# The longest allowed simulation time step, in seconds.
_DEFAULT_MAX_SIM_STEP = 1./5.


def _get_default_action(action_spec):
  """Generates an action to apply to the environment if there is no agent.

  * For action dimensions that are closed intervals this will be the midpoint.
  * For left-open or right-open intervals this will be the maximum or the
    minimum respectively.
  * For unbounded intervals this will be zero.

  Args:
    action_spec: An instance of `BoundedArraySpec` or a list or tuple
      containing these.

  Returns:
    A numpy array of actions if `action_spec` is a single `BoundedArraySpec`, or
    a tuple of such arrays if `action_spec` is a list or tuple.
  """
  if isinstance(action_spec, (list, tuple)):
    return tuple(_get_default_action(spec) for spec in action_spec)
  elif isinstance(action_spec, collections.abc.MutableMapping):
    # Clones the Mapping, preserving type and key order.
    result = copy.copy(action_spec)

    for key, value in action_spec.items():
      result[key] = _get_default_action(value)

    return result

  minimum = np.broadcast_to(action_spec.minimum, action_spec.shape)
  maximum = np.broadcast_to(action_spec.maximum, action_spec.shape)
  left_bounded = np.isfinite(minimum)
  right_bounded = np.isfinite(maximum)
  action = np.select(
      condlist=[left_bounded & right_bounded, left_bounded, right_bounded],
      choicelist=[0.5 * (minimum + maximum), minimum, maximum],
      default=0.)
  action = action.astype(action_spec.dtype, copy=False)
  action.flags.writeable = False
  return action


class State(enum.Enum):
  """State of the Runtime class."""
  START = 0
  RUNNING = 1
  STOP = 2
  STOPPED = 3
  RESTARTING = 4


class Runtime:
  """Base Runtime class.

  Attributes:
    simulation_time_budget: Float value, how much time can be spent on physics
      simulation every frame, in seconds.
    on_episode_begin: An observable subject, an instance of util.QuietSet.
      It contains argumentless callables, invoked, when a new episode begins.
    on_error: An observable subject, an instance of util.QuietSet. It contains
      single argument callables, invoked, when the environment or the agent
      throw an error.
    on_physics_changed: An observable subject, an instance of util.QuietSet.
      During episode restarts, the underlying physics instance may change. If
      you are interested in learning about those changes, attach a listener
      using the += operator. The listener should be a callable with no required
      arguments.
  """

  def __init__(self, environment, policy=None):
    """Instance initializer.

    Args:
      environment: An instance of dm_control.rl.control.Environment.
      policy: Either a callable that accepts a `TimeStep` and returns a numpy
        array of actions conforming to `environment.action_spec()`, or None, in
        which case a default action will be generated for each environment step.
    """
    self.on_error = util.QuietSet()
    self.on_episode_begin = util.QuietSet()
    self.simulation_time_budget = _DEFAULT_MAX_SIM_STEP

    self._state = State.START
    self._simulation_timer = util.Timer()
    self._tracked_simulation_time = 0.0
    self._error_logger = util.ErrorLogger(self.on_error)

    self._env = environment
    self._policy = policy
    self._default_action = _get_default_action(environment.action_spec())
    self._time_step = None
    self._last_action = None
    self.on_physics_changed = util.QuietSet()

  def tick(self, time_elapsed, paused):
    """Advances the simulation by one frame.

    Args:
      time_elapsed: Time elapsed since the last time this method was called.
      paused: A boolean flag telling if the  simulation is paused.
    Returns:
      A boolean flag to determine if the episode has finished.
    """
    with self._simulation_timer.measure_time():
      if self._state == State.RESTARTING:
        self._state = State.START
      if self._state == State.START:
        if self._start():
          self._broadcast_episode_start()
          self._tracked_simulation_time = self.get_time()
          self._state = State.RUNNING
        else:
          self._state = State.STOPPED
      if self._state == State.RUNNING:
        finished = self._step_simulation(time_elapsed, paused)
        if finished:
          self._state = State.STOP
      if self._state == State.STOP:
        self._state = State.STOPPED

  def _step_simulation(self, time_elapsed, paused):
    """Simulate a simulation step."""
    finished = False
    if paused:
      self._step_paused()
    else:
      step_duration = min(time_elapsed, self.simulation_time_budget)
      actual_simulation_time = self.get_time()
      if self._tracked_simulation_time >= actual_simulation_time:
        end_time = actual_simulation_time + step_duration
        while not finished and self.get_time() < end_time:
          finished = self._step()
      self._tracked_simulation_time += step_duration
    return finished

  def single_step(self):
    """Performs a single step of simulation."""
    if self._state == State.RUNNING:
      finished = self._step()
      self._state = State.STOP if finished else State.RUNNING

  def stop(self):
    """Stops the runtime."""
    self._state = State.STOPPED

  def restart(self):
    """Restarts the episode, resetting environment, model, and data."""
    if self._state != State.STOPPED:
      self._state = State.RESTARTING
    else:
      self._state = State.START

  def get_time(self):
    """Elapsed simulation time."""
    return self._env.physics.data.time

  @property
  def state(self):
    """Returns the current state of the state machine.

    Returned states are values of runtime.State enum.
    """
    return self._state

  @property
  def simulation_time(self):
    """Returns the amount of time spent running the simulation."""
    return self._simulation_timer.measured_time

  @property
  def last_action(self):
    """Action passed to the environment on the last step."""
    return self._last_action

  def _broadcast_episode_start(self):
    for listener in self.on_episode_begin:
      listener()

  def _start(self):
    """Starts a new simulation episode.

    Starting a new episode may be associated with changing the physics instance.
    The method tracks that and notifies observers through 'on_physics_changed'
    subject.

    Returns:
      True if the operation was successful, False otherwise.
    """
    # NB: we check the identity of the data pointer rather than the physics
    # instance itself, since this allows us to detect when the physics has been
    # "reloaded" using one of the `reload_from_*` methods.
    old_data_ptr = self._env.physics.data.ptr

    with self._error_logger:
      self._time_step = self._env.reset()

    if self._env.physics.data.ptr is not old_data_ptr:
      for listener in self.on_physics_changed:
        listener()
    return not self._error_logger.errors_found

  def _step_paused(self):
    mujoco.mj_forward(self._env.physics.model.ptr, self._env.physics.data.ptr)

  def _step(self):
    """Generates an action and applies it to the environment.

    If a `policy` was provided, this will be invoked to generate an action to
    feed to the environment, otherwise a default action will be generated.

    Returns:
      A boolean value, True if the environment signaled the episode end, False
      if the episode is still running.
    """
    finished = True
    with self._error_logger:
      if self._policy:
        action = self._policy(self._time_step)
      else:
        action = self._default_action
      self._time_step = self._env.step(action)
      self._last_action = action
      finished = self._time_step.last()
    return finished or self._error_logger.errors_found
