"""Interface for reinforcement learning."""

from deepchem.rl.a2c import A2C  # noqa: F401
from deepchem.rl.ppo import PPO  # noqa: F401


class Environment(object):
  """An environment in which an actor performs actions to accomplish a task.

  An environment has a current state, which is represented as either a single NumPy
  array, or optionally a list of NumPy arrays.  When an action is taken, that causes
  the state to be updated.  The environment also computes a reward for each action,
  and reports when the task has been terminated (meaning that no more actions may
  be taken).

  Two types of actions are supported.  For environments with discrete action spaces,
  the action is an integer specifying the index of the action to perform (out of a
  fixed list of possible actions).  For environments with continuous action spaces,
  the action is a NumPy array.

  Environment objects should be written to support pickle and deepcopy operations.
  Many algorithms involve creating multiple copies of the Environment, possibly
  running in different processes or even on different computers.
  """

  def __init__(self,
               state_shape,
               n_actions=None,
               state_dtype=None,
               action_shape=None):
    """Subclasses should call the superclass constructor in addition to doing their own initialization.

    A value should be provided for either n_actions (for discrete action spaces)
    or action_shape (for continuous action spaces), but not both.

    Parameters
    ----------
    state_shape: tuple or list of tuples
      the shape(s) of the array(s) making up the state
    n_actions: int
      the number of discrete actions that can be performed.  If the action space
      is continuous, this should be None.
    state_dtype: dtype or list of dtypes
      the type(s) of the array(s) making up the state.  If this is None, all
      arrays are assumed to be float32.
    action_shape: tuple
      the shape of the array describing an action.  If the action space
      is discrete, this should be none.
    """
    self._state_shape = state_shape
    self._n_actions = n_actions
    self._action_shape = action_shape
    self._state = None
    self._terminated = None
    if state_dtype is None:
      # Assume all arrays are float32.
      import numpy
      try:
        from collections.abc import Sequence as SequenceCollection
      except:
        from collections import Sequence as SequenceCollection
      if isinstance(state_shape[0], SequenceCollection):
        self._state_dtype = [numpy.float32] * len(state_shape)
      else:
        self._state_dtype = numpy.float32
    else:
      self._state_dtype = state_dtype

  @property
  def state(self):
    """The current state of the environment, represented as either a NumPy array or list of arrays.

    If reset() has not yet been called at least once, this is undefined.
    """
    return self._state

  @property
  def terminated(self):
    """Whether the task has reached its end.

    If reset() has not yet been called at least once, this is undefined.
    """
    return self._terminated

  @property
  def state_shape(self):
    """The shape of the arrays that describe a state.

    If the state is a single array, this returns a tuple giving the shape of that array.
    If the state is a list of arrays, this returns a list of tuples where each tuple is
    the shape of one array.
    """
    return self._state_shape

  @property
  def state_dtype(self):
    """The dtypes of the arrays that describe a state.

    If the state is a single array, this returns the dtype of that array.  If the state
    is a list of arrays, this returns a list containing the dtypes of the arrays.
    """
    return self._state_dtype

  @property
  def n_actions(self):
    """The number of possible actions that can be performed in this Environment.

    If the environment uses a continuous action space, this returns None.
    """
    return self._n_actions

  @property
  def action_shape(self):
    """The expected shape of NumPy arrays representing actions.

    If the environment uses a discrete action space, this returns None.
    """
    return self._action_shape

  def reset(self):
    """Initialize the environment in preparation for doing calculations with it.

    This must be called before calling step() or querying the state.  You can call it
    again later to reset the environment back to its original state.
    """
    raise NotImplementedError("Subclasses must implement this")

  def step(self, action):
    """Take a time step by performing an action.

    This causes the "state" and "terminated" properties to be updated.

    Parameters
    ----------
    action: object
      an object describing the action to take

    Returns
    -------
    the reward earned by taking the action, represented as a floating point number
    (higher values are better)
    """
    raise NotImplementedError("Subclasses must implement this")


class GymEnvironment(Environment):
  """This is a convenience class for working with environments from OpenAI Gym."""

  def __init__(self, name):
    """Create an Environment wrapping the OpenAI Gym environment with a specified name."""
    import gym
    self.env = gym.make(name)
    self.name = name
    space = self.env.action_space
    if 'n' in dir(space):
      super(GymEnvironment, self).__init__(self.env.observation_space.shape,
                                           space.n)
    else:
      super(GymEnvironment, self).__init__(
          self.env.observation_space.shape, action_shape=space.shape)

  def reset(self):
    self._state = self.env.reset()
    self._terminated = False

  def step(self, action):
    self._state, reward, self._terminated, info = self.env.step(action)
    return reward

  def __deepcopy__(self, memo):
    return GymEnvironment(self.name)


class Policy(object):
  """A policy for taking actions within an environment.

  A policy is defined by a tf.keras.Model that takes the current state as input
  and performs the necessary calculations.  There are many algorithms for
  reinforcement learning, and they differ in what values they require a policy to
  compute.  That makes it impossible to define a single interface allowing any
  policy to be optimized with any algorithm.  Instead, this interface just tries
  to be as flexible and generic as possible.  Each algorithm must document what
  values it expects the model to output.

  Special handling is needed for models that include recurrent layers.  In that
  case, the model has its own internal state which the learning algorithm must
  be able to specify and query.  To support this, the Policy must do three things:

  1. The Model must take additional inputs that specify the initial states of
     all its recurrent layers.  These will be appended to the list of arrays
     specifying the environment state.

  2. The Model must also return the final states of all its recurrent layers as
     outputs.

  3. The constructor argument rnn_initial_states must be specified to define
     the states to use for the Model's recurrent layers at the start of a new
     rollout.

  Policy objects should be written to support pickling.  Many algorithms involve
  creating multiple copies of the Policy, possibly running in different processes
  or even on different computers.
  """

  def __init__(self, output_names, rnn_initial_states=[]):
    """Subclasses should call the superclass constructor in addition to doing
    their own initialization.

    Parameters
    ----------
    output_names: list of strings
      the names of the Model's outputs, in order.  It is up to each reinforcement
      learning algorithm to document what outputs it expects policies to compute.
      Outputs that return the final states of recurrent layers should have the
      name 'rnn_state'.
    rnn_initial_states: list of NumPy arrays
      the initial states of the Model's recurrent layers at the start of a new
      rollout
    """
    self.output_names = output_names
    self.rnn_initial_states = rnn_initial_states

  def create_model(self, **kwargs):
    """Construct and return a tf.keras.Model that computes the policy.

    The inputs to the model consist of the arrays representing the current state
    of the environment, followed by the initial states for all recurrent layers.
    Depending on the algorithm being used, other inputs might get passed as
    well.  It is up to each algorithm to document that.
    """
    raise NotImplementedError("Subclasses must implement this")
