# Copyright 2019 DeepMind Technologies Limited. 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.
# ==============================================================================
"""JAX functions for state value and action-value learning.

Value functions estimate the expected return (discounted sum of rewards) that
can be collected by an agent under a given policy of behaviour. This subpackage
implements a number of functions for value learning in discrete scalar action
spaces. Actions are assumed to be represented as indices in the range `[0, A)`
where `A` is the number of distinct actions.
"""
from typing import Union
import chex
import jax
import jax.numpy as jnp
from rlax._src import base
from rlax._src import clipping
from rlax._src import distributions
from rlax._src import multistep

Array = chex.Array
Numeric = chex.Numeric


def td_learning(
    v_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    v_t: Numeric,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the TD-learning temporal difference error.

  See "Learning to Predict by the Methods of Temporal Differences" by Sutton.
  (https://link.springer.com/article/10.1023/A:1022633531479).

  Args:
    v_tm1: state values at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    v_t: state values at time t.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.
  Returns:
    TD-learning temporal difference error.
  """

  chex.assert_rank([v_tm1, r_t, discount_t, v_t], 0)
  chex.assert_type([v_tm1, r_t, discount_t, v_t], float)

  target_tm1 = r_t + discount_t * v_t
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - v_tm1


def td_lambda(
    v_tm1: Array,
    r_t: Array,
    discount_t: Array,
    v_t: Array,
    lambda_: Numeric,
    stop_target_gradients: bool = True,
) -> Array:
  """Calculates the TD(lambda) temporal difference error.

  See "Reinforcement Learning: An Introduction" by Sutton and Barto.
  (http://incompleteideas.net/book/ebook/node74.html).

  Args:
    v_tm1: sequence of state values at time t-1.
    r_t: sequence of rewards at time t.
    discount_t: sequence of discounts at time t.
    v_t: sequence of state values at time t.
    lambda_: mixing parameter lambda, either a scalar or a sequence.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.
  Returns:
    TD(lambda) temporal difference error.
  """
  chex.assert_rank([v_tm1, r_t, discount_t, v_t, lambda_], [1, 1, 1, 1, {0, 1}])
  chex.assert_type([v_tm1, r_t, discount_t, v_t, lambda_], float)

  target_tm1 = multistep.lambda_returns(r_t, discount_t, v_t, lambda_)
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - v_tm1


def sarsa(
    q_tm1: Array,
    a_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    q_t: Array,
    a_t: Numeric,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the SARSA temporal difference error.

  See "Reinforcement Learning: An Introduction" by Sutton and Barto.
  (http://incompleteideas.net/book/ebook/node64.html.)

  Args:
    q_tm1: Q-values at time t-1.
    a_tm1: action index at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    q_t: Q-values at time t.
    a_t: action index at time t.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.
  Returns:
    SARSA temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t, a_t],
                   [1, 0, 0, 0, 1, 0])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t, a_t],
                   [float, int, float, float, float, int])

  target_tm1 = r_t + discount_t * q_t[a_t]
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - q_tm1[a_tm1]


def expected_sarsa(
    q_tm1: Array,
    a_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    q_t: Array,
    probs_a_t: Array,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the expected SARSA (SARSE) temporal difference error.

  See "A Theoretical and Empirical Analysis of Expected Sarsa" by Seijen,
  van Hasselt, Whiteson et al.
  (http://www.cs.ox.ac.uk/people/shimon.whiteson/pubs/vanseijenadprl09.pdf).

  Args:
    q_tm1: Q-values at time t-1.
    a_tm1: action index at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    q_t: Q-values at time t.
    probs_a_t: action probabilities at time t.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.

  Returns:
    Expected SARSA temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t, probs_a_t],
                   [1, 0, 0, 0, 1, 1])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t, probs_a_t],
                   [float, int, float, float, float, float])

  target_tm1 = r_t + discount_t * jnp.dot(q_t, probs_a_t)
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - q_tm1[a_tm1]


def sarsa_lambda(
    q_tm1: Array,
    a_tm1: Array,
    r_t: Array,
    discount_t: Array,
    q_t: Array,
    a_t: Array,
    lambda_: Numeric,
    stop_target_gradients: bool = True,
) -> Array:
  """Calculates the SARSA(lambda) temporal difference error.

  See "Reinforcement Learning: An Introduction" by Sutton and Barto.
  (http://incompleteideas.net/book/ebook/node77.html).

  Args:
    q_tm1: sequence of Q-values at time t-1.
    a_tm1: sequence of action indices at time t-1.
    r_t: sequence of rewards at time t.
    discount_t: sequence of discounts at time t.
    q_t: sequence of Q-values at time t.
    a_t: sequence of action indices at time t.
    lambda_: mixing parameter lambda, either a scalar or a sequence.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.

  Returns:
    SARSA(lambda) temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t, a_t, lambda_],
                   [2, 1, 1, 1, 2, 1, {0, 1}])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t, a_t, lambda_],
                   [float, int, float, float, float, int, float])

  qa_tm1 = base.batched_index(q_tm1, a_tm1)
  qa_t = base.batched_index(q_t, a_t)
  target_tm1 = multistep.lambda_returns(r_t, discount_t, qa_t, lambda_)
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - qa_tm1


def q_learning(
    q_tm1: Array,
    a_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    q_t: Array,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the Q-learning temporal difference error.

  See "Reinforcement Learning: An Introduction" by Sutton and Barto.
  (http://incompleteideas.net/book/ebook/node65.html).

  Args:
    q_tm1: Q-values at time t-1.
    a_tm1: action index at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    q_t: Q-values at time t.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.

  Returns:
    Q-learning temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t], [1, 0, 0, 0, 1])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t],
                   [float, int, float, float, float])

  target_tm1 = r_t + discount_t * jnp.max(q_t)
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - q_tm1[a_tm1]


def double_q_learning(
    q_tm1: Array,
    a_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    q_t_value: Array,
    q_t_selector: Array,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the double Q-learning temporal difference error.

  See "Double Q-learning" by van Hasselt.
  (https://papers.nips.cc/paper/3964-double-q-learning.pdf).

  Args:
    q_tm1: Q-values at time t-1.
    a_tm1: action index at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    q_t_value: Q-values at time t.
    q_t_selector: selector Q-values at time t.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.

  Returns:
    Double Q-learning temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t_value, q_t_selector],
                   [1, 0, 0, 0, 1, 1])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t_value, q_t_selector],
                   [float, int, float, float, float, float])

  target_tm1 = r_t + discount_t * q_t_value[q_t_selector.argmax()]
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - q_tm1[a_tm1]


def persistent_q_learning(
    q_tm1: Array,
    a_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    q_t: Array,
    action_gap_scale: float,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the persistent Q-learning temporal difference error.

  See "Increasing the Action Gap: New Operators for Reinforcement Learning"
  by Bellemare, Ostrovski, Guez et al. (https://arxiv.org/abs/1512.04860).

  Args:
    q_tm1: Q-values at time t-1.
    a_tm1: action index at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    q_t: Q-values at time t.
    action_gap_scale: coefficient in [0, 1] for scaling the action gap term.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.

  Returns:
    Persistent Q-learning temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t], [1, 0, 0, 0, 1])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t],
                   [float, int, float, float, float])

  corrected_q_t = (
      (1. - action_gap_scale) * jnp.max(q_t)
      + action_gap_scale * q_t[a_tm1]
  )
  target_tm1 = r_t + discount_t * corrected_q_t
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - q_tm1[a_tm1]


def qv_learning(
    q_tm1: Array,
    a_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    v_t: Numeric,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the QV-learning temporal difference error.

  See "Two Novel On-policy Reinforcement Learning Algorithms based on
  TD(lambda)-methods" by Wiering and van Hasselt
  (https://ieeexplore.ieee.org/abstract/document/4220845.)

  Args:
    q_tm1: Q-values at time t-1.
    a_tm1: action index at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    v_t: state values at time t.
    stop_target_gradients: bool indicating whether or not to apply stop gradient
      to targets.

  Returns:
    QV-learning temporal difference error.
  """
  chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, v_t], [1, 0, 0, 0, 0])
  chex.assert_type([q_tm1, a_tm1, r_t, discount_t, v_t],
                   [float, int, float, float, float])

  target_tm1 = r_t + discount_t * v_t
  target_tm1 = jax.lax.select(stop_target_gradients,
                              jax.lax.stop_gradient(target_tm1), target_tm1)
  return target_tm1 - q_tm1[a_tm1]


def qv_max(
    v_tm1: Numeric,
    r_t: Numeric,
    discount_t: Numeric,
    q_t: Array,
    stop_target_gradients: bool = True,
) -> Numeric:
  """Calculates the QVMAX temporal difference error.

  See "The QV Family Compared to Other Reinforcement Learning Algorithms" by
  Wiering and van Hasselt (2009).
  (http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.713.1931)

  Args:
    v_tm1: state values at time t-1.
    r_t: reward at time t.
    discount_t: discount at time t.
    q_t: Q-values at time t.