# Copyright 2018-2019 The dm_control Authors.
# Copyright 2023 The RoboPianist 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 typing import Optional

import dm_env
import mujoco
import numpy as np
from dm_control.mujoco.wrapper import mjbindings

from robopianist.models.piano import midi_module
from robopianist.music import synthesizer
from robopianist.viewer import util

mjlib = mjbindings.mjlib

# The longest allowed simulation time step, in seconds.
_DEFAULT_MAX_SIM_STEP = 1.0 / 5.0


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.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) -> 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._environment = environment
        self._policy = policy
        self._default_action = _get_default_action(environment.action_spec())
        self._time_step: Optional[dm_env.TimeStep] = None
        self._last_action = None
        self.on_physics_changed = util.QuietSet()

        # Initialize the synthesizer.
        self._synthesizer = synthesizer.Synthesizer()
        self._synthesizer.start()
        self._midi_module = environment.task.piano.midi_module
        assert isinstance(self._midi_module, midi_module.MidiModule)
        self._midi_module.register_synth_note_on_callback(self._synthesizer.note_on)
        self._midi_module.register_synth_note_off_callback(self._synthesizer.note_off)
        self._midi_module.register_synth_sustain_on_callback(
            self._synthesizer.sustain_on
        )
        self._midi_module.register_synth_sustain_off_callback(
            self._synthesizer.sustain_off
        )

    def tick(self, time_elapsed, paused, muted):
        """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.
            muted: A boolean flag telling if the sound is muted.

        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:
                self._synthesizer.mute(muted)
                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
        if paused or finished:
            if self._synthesizer.sustained:
                self._synthesizer.sustain_off()
            self._synthesizer.all_notes_off()
        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._environment.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._environment.physics.data.ptr

        with self._error_logger:
            self._synthesizer.all_sounds_off()
            self._time_step = self._environment.reset()

        if self._environment.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._environment.physics.model.ptr, self._environment.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._environment.step(action)
            self._last_action = action
            finished = self._time_step.last()
        return finished or self._error_logger.errors_found

    @property
    def environment(self):
        return self._environment

    def __del__(self):
        self._synthesizer.stop()
