# Copyright (c) 2020 Horizon Robotics and ALF Contributors. All Rights Reserved.
#
# 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.
"""Step a single env in a separate process for lock free paralellism.

Adapted from TF-Agents Environment API as seen in:
    https://github.com/tensorflow/agents/blob/master/tf_agents/environments/parallel_py_environment.py
"""

from absl import logging
import atexit
from enum import Enum
import multiprocessing
import numpy as np
import sys
import torch
import traceback

import alf
from alf.data_structures import TimeStep
import alf.nest as nest


class _MessageType(Enum):
    """Message types for communication via the pipe.

    The ProcessEnvironment uses pipe to perform IPC, where each of the message
    has a message type. This Enum provides all the available message types.
    """
    READY = 1
    ACCESS = 2
    CALL = 3
    RESULT = 4
    EXCEPTION = 5
    CLOSE = 6


def _worker(conn, env_constructor, env_id=None, flatten=False):
    """The process waits for actions and sends back environment results.

    Args:
        conn (multiprocessing.connection): Connection for communication to the main process.
        env_constructor (Callable): callable environment creator.
        flatten (bool): whether to assume flattened actions and time_steps
          during communication to avoid overhead.

    Raises:
        KeyError: When receiving a message of unknown type.
    """
    try:
        alf.set_default_device("cpu")
        env = env_constructor(env_id=env_id)
        action_spec = env.action_spec()
        conn.send(_MessageType.READY)  # Ready.
        while True:
            try:
                # Only block for short times to have keyboard exceptions be raised.
                if not conn.poll(0.1):
                    continue
                message, payload = conn.recv()
            except (EOFError, KeyboardInterrupt):
                break
            if message == _MessageType.ACCESS:
                name = payload
                result = getattr(env, name)
                conn.send((_MessageType.RESULT, result))
                continue
            if message == _MessageType.CALL:
                name, args, kwargs = payload
                if flatten and name == 'step':
                    args = [nest.pack_sequence_as(action_spec, args[0])]
                result = getattr(env, name)(*args, **kwargs)
                if flatten and name in ['step', 'reset']:
                    result = nest.flatten(result)
                    assert all([
                        not isinstance(x, torch.Tensor) for x in result
                    ]), ("Tensor result is not allowed: %s" % name)
                conn.send((_MessageType.RESULT, result))
                continue
            if message == _MessageType.CLOSE:
                assert payload is None
                env.close()
                break
            raise KeyError(
                'Received message of unknown type {}'.format(message))
    except KeyboardInterrupt:
        # When worker receives interruption from keyboard (i.e. Ctrl-C), notify
        # the parent process to shut down quietly by sending the CLOSE message.
        #
        # This is to avoid sometimes tens of environment processes panicking
        # simultaneously.
        conn.send((_MessageType.CLOSE, None))
    except Exception:  # pylint: disable=broad-except
        etype, evalue, tb = sys.exc_info()
        stacktrace = ''.join(traceback.format_exception(etype, evalue, tb))
        message = 'Error in environment process: {}'.format(stacktrace)
        logging.error(message)
        conn.send((_MessageType.EXCEPTION, stacktrace))
    finally:
        conn.close()


class ProcessEnvironment(object):
    def __init__(self, env_constructor, env_id=None, flatten=False):
        """Step environment in a separate process for lock free paralellism.

        The environment is created in an external process by calling the provided
        callable. This can be an environment class, or a function creating the
        environment and potentially wrapping it. The returned environment should
        not access global variables.

        Args:
            env_constructor (Callable): callable environment creator.
            env_id (torch.int32): ID of the the env
            flatten (bool): whether to assume flattened actions and time_steps
                during communication to avoid overhead.

        Attributes:
            observation_spec: The cached observation spec of the environment.
            action_spec: The cached action spec of the environment.
            time_step_spec: The cached time step spec of the environment.
        """
        self._env_constructor = env_constructor
        self._flatten = flatten
        self._env_id = env_id
        self._observation_spec = None
        self._action_spec = None
        self._reward_spec = None
        self._time_step_spec = None
        self._env_info_spec = None

    def start(self, wait_to_start=True):
        """Start the process.

        Args:
            wait_to_start (bool): Whether the call should wait for an env initialization.
        """
        # The following context made sure that the newly created child process
        # (for environment) is started using the "fork" start method.
        #
        # This is to prevent multiprocessing from accidentally creating the
        # child process with the "spawn" start method. Using "fork" start method
        # is required here because we would like to have the child process
        # inherit the alf configurations from the parent process, so that such
        # confiurations are effective for the to-be-created environments in the
        # child process.
        mp_ctx = multiprocessing.get_context('fork')
        self._conn, conn = mp_ctx.Pipe()
        self._process = mp_ctx.Process(
            target=_worker,
            args=(conn, self._env_constructor, self._env_id, self._flatten))
        atexit.register(self.close)
        self._process.start()
        if wait_to_start:
            self.wait_start()

    def wait_start(self):
        """Wait for the started process to finish initialization."""
        result = self._conn.recv()
        if isinstance(result, Exception):
            self._conn.close()
            self._process.join(5)
            raise result
        assert result == _MessageType.READY, result

    def env_info_spec(self):
        if not self._env_info_spec:
            self._env_info_spec = self.call('env_info_spec')()
        return self._env_info_spec

    def observation_spec(self):
        if not self._observation_spec:
            self._observation_spec = self.call('observation_spec')()
        return self._observation_spec

    def action_spec(self):
        if not self._action_spec:
            self._action_spec = self.call('action_spec')()
        return self._action_spec

    def reward_spec(self):
        if not self._reward_spec:
            self._reward_spec = self.call('reward_spec')()
        return self._reward_spec

    def time_step_spec(self):
        if not self._time_step_spec:
            self._time_step_spec = self.call('time_step_spec')()
        return self._time_step_spec

    def __getattr__(self, name):
        """Request an attribute from the environment.

        Note that this involves communication with the external process, so it can
        be slow.

        Args:
            name (str): Attribute to access.

        Returns:
            Value of the attribute.
        """
        self._conn.send((_MessageType.ACCESS, name))
        return self._receive()

    def call(self, name, *args, **kwargs):
        """Asynchronously call a method of the external environment.

        Args:
            name (str): Name of the method to call.
            *args: Positional arguments to forward to the method.
            **kwargs: Keyword arguments to forward to the method.

        Returns:
            Promise object that blocks and provides the return value when called.
        """
        payload = name, args, kwargs
        self._conn.send((_MessageType.CALL, payload))
        return self._receive

    def close(self):
        """Send a close message to the external process and join it."""
        try:
            self._conn.send((_MessageType.CLOSE, None))
            self._conn.close()
        except IOError:
            # The connection was already closed.
            pass
        self._process.join(5)

    def step(self, action, blocking=True):
        """Step the environment.

        Args:
            action (nested tensors): The action to apply to the environment.
            blocking (bool): Whether to wait for the result.

        Returns:
            time step when blocking, otherwise callable that returns the time step.
        """
        promise = self.call('step', action)
        if blocking:
            return promise()
        else:
            return promise

    def reset(self, blocking=True):
        """Reset the environment.

        Args:
            blocking (bool): Whether to wait for the result.

        Returns:
            New observation when blocking, otherwise callable that returns the new
            observation.
        """
        promise = self.call('reset')
        if blocking:
            return promise()
        else:
            return promise

    def _receive(self):
        """Wait for a message from the worker process and return its payload.

        Raises:
            Exception: An exception was raised inside the worker process.
            KeyError: The reveived message is of an unknown type.

        Returns:
            Payload object of the message.
        """
        message, payload = self._conn.recv()

        # Re-raise exceptions in the main process.
        if message == _MessageType.EXCEPTION:
            stacktrace = payload
            raise Exception(stacktrace)
        elif message == _MessageType.RESULT:
            return payload
        elif message == _MessageType.CLOSE:
            # When notified that the child process is going to shut down, do not
            # panic and handle it quietly.
            return None
        self.close()
        raise KeyError(
            'Received message of unexpected type {}'.format(message))

    def render(self, mode='human'):
        """Render the environment.

        Args:
            mode (str): One of ['rgb_array', 'human']. Renders to an numpy array, or brings
                up a window where the environment can be visualized.
        Returns:
            An ndarray of shape [width, height, 3] denoting an RGB image if mode is
            `rgb_array`. Otherwise return nothing and render directly to a display
            window.
        Raises:
            NotImplementedError: If the environment does not support rendering.
        """
        return self.call('render', mode)()
