# Copyright 2021 DeepMind Technologies Limited
#
# 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.
# ==============================================================================

"""Policies are functions mapping contexts to actions.

Policies are described in
"Kuzborskij, I., Vernade, C., Gyorgy, A., & Szepesvári, C. (2021, March).
Confident off-policy evaluation and selection through self-normalized importance
weighting. In International Conference on Artificial Intelligence and Statistics
(pp. 640-648). PMLR.".

class SoftmaxDataPolicy is a mock-up policy which can hold either training
sample or a testing sample (each of which consists of context and labels).
When either set of contexts is passed to the policy (get_probs(...))
it returns action probabilities associated with those contexts.
Note that this is a mock-up policy, so only one of the two samples is supported.

class SoftmaxGAPolicy implements a softmax policy with linear parameterized
potential, where parameters are fitted by the gradient ascent maximizing
either importance weighted or self-normalized importance weighted estimator.
"""
import abc
import enum
import math
from typing import Sequence, NamedTuple

from absl import logging
import jax
from jax import numpy as jnp
from jax import scipy as jsc
import numpy as np
import scipy
import sklearn.preprocessing as skl_prep

from offpolicy_selection_eslb.utils import sample_from_simplices_m_times


class Query(NamedTuple):
  """Actions generated by a (randomized) policy when given a set of contexts.

  Attributes:
    actions: n-times-1 Array -- chosen (sampled) actions
    probabilities: n-times-1 Array -- corresponding probabilities
  """

  actions: np.ndarray
  probabilities: np.ndarray


def log_vhat_importance_weighting(
    parameters: np.ndarray,
    temperature: float,
    contexts: np.ndarray,
    actions: np.ndarray,
    rewards: np.ndarray,
    b_prob: np.ndarray,
) -> np.ndarray:
  """Returns the log of importance weighted estimator.

  Returns the log of importance weighted estimator where each
  importance weight is computed w.r.t. the softmax target policy defined
  w.r.t. a linear model as defined in the description of a class.

  Args:
    parameters: Parameters of the linear model of a target policy.
    temperature: Positive float controlling the temperature of a Softmax
      policy.
    contexts: Array of contexts (n-times-d, d=data dim., n=sample size).
    actions: Actions (integers).
    rewards: Rewards (float).
    b_prob: Probabilities corresponding to (context, action) pairs
      according to the behavior policy.
  Returns: The logarithm of importance-weighted estimate.
  """
  n, _ = contexts.shape
  v = (1.0 / temperature) * contexts.dot(parameters)
  pot = (1.0 / temperature) * (contexts *
                               parameters[:, actions].T).sum(axis=1)

  a = jnp.log(rewards / (n * b_prob)) - jsc.special.logsumexp(v, axis=1)
  rs = jsc.special.logsumexp(pot + a, axis=0)

  return rs


def log_vhat_sn_importance_weighting(
    parameters: np.ndarray,
    temperature: float,
    contexts: np.ndarray,
    actions: np.ndarray,
    rewards: np.ndarray,
    b_prob: np.ndarray,
) -> np.ndarray:
  """Returns a log of self-normalized (SN) importance weighted estimator.

  Returns a log of (SN) importance weighted estimator where each
  importance weight is computed w.r.t. the softmax target policy defined
  w.r.t. a linear model as defined in the description of a class.

  Args:
    parameters: Parameters of the linear model of a target policy.
    temperature: Positive float controlling the temperature of a Softmax
      policy.
    contexts: Array of contexts (n-times-d, d=data dim., n=sample size).
    actions: Actions (integers).
    rewards: Rewards (float).
    b_prob: Probabilities corresponding to (context, action) pairs
      according to the behavior policy.
  Returns: The logarithm of SN importance-weighted estimate.
  """
  v = (1.0 / temperature) * contexts.dot(parameters)
  pot = (1.0 / temperature) * (contexts *
                               parameters[:, actions].T).sum(axis=1)

  a = jnp.log(rewards / b_prob) - jsc.special.logsumexp(v, axis=1)
  ln_numer = jsc.special.logsumexp(pot + a, axis=0)

  a = -jnp.log(b_prob) - jsc.special.logsumexp(v, axis=1)
  ln_denom = jsc.special.logsumexp(pot + a, axis=0)

  return ln_numer - ln_denom


class Policy(abc.ABC):
  """A Policy samples actions given contexts.
  """

  @abc.abstractmethod
  def query(self, contexts: np.ndarray) -> Query:
    """Returns actions and their probs sampled by Policy given the contexts.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size), which
        are either training or testing contexts provided during the
        initialization.
    Returns: A Tuple of arrays of actions (int) and corresponding probs (float)
    """

  @abc.abstractmethod
  def get_probs(self, contexts: np.ndarray) -> np.ndarray:
    """Returns probability distribution over actions for each context.

    The softmax policy is defined as a probability vector
      exp(alt_bin_labels / temp) / sum(exp(alt_bin_labels / temp))
      where temp is a temperature of a policy and
      alt_bin_labels is a binary encoding of labels altered by alter_labels(...)

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size), which
        are either training or testing contexts provided during the
        initialization.
    Returns:  Array of probabilities according to the policy, where K
      is the number of actions (size n-times-K).

    Raises:
      NotImplementedError: when contexts is not training or testing contexts
    """


class TrainedPolicyObjType(enum.Enum):
  """Softmax gradient ascent fitted policy types with Objective function.

  TrainedPolicyObjType.IW = importance-weighted estimator.
  TrainedPolicyObjType.SNIW = self-normalized importance-weighted.
  """
  IW = "IW"
  SNIW = "SNIW"

  def __str__(self):
    return str(self.value)


class SoftmaxDataPolicy(Policy):
  """Memorization policy (using true labels).

  This object can hold either training sample or a testing sample
  (each of which consists of context and labels).
  When either set of contexts is passed to the policy (get_probs(...))
  it returns action probabilities associated with those contexts.
  Note that this is a mock-up policy, so only one of the two samples is
  supported.


  Attributes:
    action_set: A list of unique integer actions.
    train_contexts: A n-times-d array of training contexts(d=data dim., n=sample
      size).
    train_labels: A n-array of training labels.
    test_contexts: A n-times-d array of training contexts(d=data dim., n'=sample
      size).
    test_labels: A n'-array of training labels.
    temperature: A positive float controlling the temp. of a Softmax policy.
    faulty_actions: A list of labels where the behavior policy makes mistakes.
    rand: Random state of numpy.random.RandomState type.
  """

  def __init__(
      self,
      train_contexts: np.ndarray,
      train_labels: np.ndarray,
      test_contexts: np.ndarray,
      test_labels: np.ndarray,
      action_set: Sequence[int],
      temperature: float,
      faulty_actions: Sequence[int],
  ):
    """Constructs a Policy.


    Args:
      train_contexts:  Array of training contexts (n-times-d, d=data dim.,
        n=sample size).
      train_labels: Array of training labels (size n).
      test_contexts: Array of training contexts (n-times-d, d=data dim.,
        n'=sample size).
      test_labels: Array of training labels 9size n).
      action_set: List of unique integer actions.
      temperature: Positive float controlling the temperature of a Softmax
        policy.
      faulty_actions: List of labels on which the behavior policy makes
        mistakes.
    """
    self.action_set = action_set

    self.train_contexts = train_contexts
    self.train_labels = train_labels
    self.test_contexts = test_contexts
    self.test_labels = test_labels
    self.temperature = temperature
    self.faulty_actions = set(faulty_actions)
    self.reset_noise(0)

  def reset_noise(self, seed: int):
    """Resets a random state given a seed.

    Args:
      seed: Integer seed for random state
    """
    self.rand = np.random.RandomState(seed)

  def alter_labels(self, labels: np.ndarray):
    """Returns altered labels according to the self.faulty_actions spec.

    Labels are altered by shifting each label contained in self.faulty_action
    to one forward (or to 0 if we have an overflow).

    Args:
      labels: Vector of labels (size 1 by n=sample size)

    Returns:
      A vector of the same size with all entries in self.faulty_actions shifted.
    """
    num_actions = len(self.action_set)

    fault = np.zeros(len(labels))
    for i in range(len(labels)):
      if labels[i] in self.faulty_actions:
        fault[i] = 1

    return (labels + fault) % num_actions  # faulty actions get shifted by one

  def get_probs(self, contexts: np.ndarray):
    """Returns probability distribution over actions for given contexts.

    The softmax policy is defined as a probability vector
      exp(alt_bin_labels / temp) / sum(exp(alt_bin_labels / temp))
      where temp is a temperature of a policy and
      alt_bin_labels is a binary encoding of labels altered by alter_labels(...)

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size), which
        are either training or testing contexts provided during the
        initialization.
    Returns:  Array of probabilities according to the policy, where K
      is the number of actions (size n-times-K).

    Raises:
      NotImplementedError: when contexts is not training or testing contexts
    """

    # predictions get altered by internal noise :
    if contexts is self.train_contexts:
      alt_labels = self.alter_labels(self.train_labels)
    elif contexts is self.test_contexts:
      alt_labels = self.alter_labels(self.test_labels)
    else:
      raise NotImplementedError

    bin_alt_labels = skl_prep.label_binarize(
        alt_labels, classes=self.action_set)

    v = np.exp(bin_alt_labels / self.temperature)
    v = v / v.sum(axis=1)[:, np.newaxis]

    return v

  def get_probs_by_actions(self, contexts: np.ndarray, actions: np.ndarray):
    """Returns probabilities for each given action in each given context.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size), which
        are either training or testing contexts provided during the
        initialization.
      actions: Array of actions (integers) for which probabilies are
        requested.
    Returns: Probabilities according to the policy.
    """
    n = len(actions)
    all_probs = self.get_probs(contexts)
    probs = all_probs[np.arange(n), actions]
    return probs

  def query(self, contexts: np.ndarray) -> Query:
    """Returns actions and their probs sampled for the given contexts.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size), which
        are either training or testing contexts provided during the
        initialization.
    Returns: A Tuple of arrays of actions (int) and corresponding probs (float)
    """
    probs = self.get_probs(contexts)
    actions = [np.random.choice(self.action_set, p=pi) for pi in probs]

    n = probs.shape[0]
    probs_by_actions = probs[:n, actions]
    return Query(np.array(actions), np.array(probs_by_actions))

  def __str__(self):
    """Returns a string representation of a policy with parametrization."""
    return f"SoftmaxDataPolicy(τ={self.temperature}, fauly_actions=[{str(self.faulty_actions)}])"


class SoftmaxGAPolicy(Policy):
  """Softmax gradient ascent fitted policy.

  This softmax policy is defined as a probability vector
  x |-> exp(<W,x> / temp) / sum(exp(<W,x> / temp))
      where temp is a temperature of a policy and
      W is a K-times-d matrix of parameters (here K is a number of actions
      and d is a context dimension).
  Parameters W are fitted by the gradient ascent either w.r.t. the
  importance-weighted estimator or its self-normalized version.

  Attributes:
    n_actions: Number of actions.
    temperature: Positive float controlling the temp. of a Softmax policy.
    steps: Number of gradient ascent steps for fitting the policy
      step_size: step size of the gradient ascent for fitting the policy.
    obj_type: Objective type, TrainedPolicyObjType.IW = importance-weighted
      estimator TrainedPolicyObjType.SNIW = self-normalized importance-weighted
      estimator.
    parameters: Parameters of the linear model in the softmax policy
    ln_obj: Reference to a static method implementing the
      log-objective function.
  """

  def __init__(
      self,
      action_set: Sequence[int],
      temperature: float,
      steps: int = 10000,
      step_size: float = 1e-2,
      obj_type: TrainedPolicyObjType = TrainedPolicyObjType.IW,
  ):
    """Constructs a Softmax Gradient Ascent Policy.

    Args:
      action_set: List of unique integer actions.
      temperature: Positive float controlling the temperature of a Softmax
        policy.
      steps: Number of gradient ascent steps for fitting the policy.
      step_size: Step size of the gradient ascent for fitting the policy.
      obj_type: Objective type, TrainedPolicyObjType.IW = importance-weighted
        estimator TrainedPolicyObjType.SNIW = self-normalized
        importance-weighted estimator.
    """
    self.n_actions = len(action_set)
    self.temperature = temperature
    self.steps = steps
    self.step_size = step_size
    self.parameters = None

    self.obj_type = obj_type
    if obj_type == TrainedPolicyObjType.IW:
      self.ln_obj = log_vhat_importance_weighting
    elif obj_type == TrainedPolicyObjType.SNIW:
      self.ln_obj = log_vhat_sn_importance_weighting
    else:
      raise NotImplementedError

  def train(
      self,
      contexts: np.ndarray,
      actions: np.ndarray,
      rewards: np.ndarray,
      b_prob: np.ndarray,
  ):
    """Fits the softmax policy according to the chosen objective.

    Fits the softmax policy according to the objective chosen during
    initialization. The gradient ascent is run for a fixed number of
    steps and a step size (specified during initialization).
    Gradient computation is done through autodiff jax library.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size)
      actions: Actions (integers).
      rewards: Rewards (float).
      b_prob: Probabilities corresponding to (context, action) pairs
        according to the behavior policy.
    """
    contexts = jnp.array(contexts)
    actions = jnp.array(actions)
    rewards = jnp.array(rewards)
    b_prob = jnp.array(b_prob)

    _, d = contexts.shape

    grad_v = jax.jit(jax.grad(self.ln_obj))

    obj_params = (self.temperature, contexts, actions, rewards, b_prob)

    logging.debug("%s(softmax): iter\t\temp_value ", self.obj_type)
    logging.debug("%s(softmax): --------------------------------- ",
                  self.obj_type)

    def update_step_ga(_, parameters: np.ndarray):
      """Returns updated parameters after a single step of gradient ascent.

      Args:
        _: gradient ascent step
        parameters: Parameters to be updated.
      Returns: Updated parameters.
      """
      g = grad_v(parameters, *obj_params)
      parameters += self.step_size * g
      return parameters

    parameters_init = np.zeros(shape=(d, self.n_actions))
    parameters_init = jnp.array(parameters_init)

    self.parameters = jax.lax.fori_loop(0, self.steps, update_step_ga,
                                        parameters_init)

    logging.debug("%s(softmax): %d\t\t%.2f ", self.obj_type, self.steps,
                  math.exp(self.ln_obj(self.parameters, *obj_params)))

  def get_probs(self, contexts: np.ndarray):
    """Returns probability distribution over actions for the given contexts.

    The softmax policy is defined as a probability vector
      exp(<W,x> / temp) / sum(exp(<W,x> / temp))
      where temp is a temperature of a policy and
      W is a K-times-d matrix of parameters (here K is a number of actions
      and d is a context dimension) fitted by gradient ascent.

    Args:
      contexts:  Array of contexts (n-times-d, d=data dim., n=sample size).
    Returns: Array of probabilities according to the policy.
    """
    return np.exp(self.get_logprobs(contexts))

  def get_probs_by_actions(self, contexts, actions):
    """Returns probability for each given action in each given context.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size).
      actions: Actions (integers) for which probabilies are requested.
    Returns: Probabilities according to the policy.
    """
    return np.exp(self.get_logprobs_by_actions(contexts, actions))

  def get_logprobs(self, contexts: np.ndarray):
    """Returns log-probabilities over actions for each given context.

    The softmax policy is defined as a probability vector
      log(exp(<W,x> / temp) / sum(exp(<W,x> / temp)))
      where temp is a temperature of a policy and
      W is a K-times-d matrix of parameters (here K is a number of actions
      and d is a context dimension) fitted by gradient ascent.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size).
    Returns: Array of log-probabilities according to the policy (n-times-K).
    """
    v = (1.0 / self.temperature) * contexts.dot(self.parameters)
    logprob = v - np.expand_dims(scipy.special.logsumexp(v, axis=1), axis=1)
    return logprob

  def get_logprobs_by_actions(self, contexts, actions):
    """Returns log-probabilities for each given action and context.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size).
      actions: Actions (integers) for which probabilies are requested.
    Returns: Log-probabilities according to the policy.
    """
    v = (1.0 / self.temperature) * contexts.dot(self.parameters)
    pot = (1.0 / self.temperature) * (contexts *
                                      self.parameters[:, actions].T).sum(axis=1)
    logprob = pot - scipy.special.logsumexp(v, axis=1)
    return logprob

  def query(self, contexts: np.ndarray) -> Query:
    """Returns actions and their probs sampled by the policy given the contexts.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size), which
        are either training or testing contexts provided during the
        initialization.
    Returns: Array integer actions and array of corresponding probabilities.
    """
    n = len(contexts)
    probs = np.zeros((self.n_actions, n))

    actions = []
    sample_probs = []

    probs = self.get_probs(contexts)
    actions = sample_from_simplices_m_times(p=probs, m=1).squeeze()
    sample_probs = probs[:n, actions]

    return Query(np.array(actions), np.array(sample_probs))

  def query_many_times(self, contexts: np.ndarray, m_times: int):
    """Returns m_times actions sampled according to Policy for each context.

    Samples actions m_times times efficiently.

    Args:
      contexts: Array of contexts (n-times-d, d=data dim., n=sample size).
      m_times: Number of times to repeat the query.
    Returns: Array of integer actions (n-times-m_times) and n-array of
      corresponding probabilities.
    """
    n = len(contexts)
    k = self.n_actions
    probs = np.zeros((k, n))

    actions = []
    sample_probs = []

    probs = self.get_probs(contexts)
    actions = sample_from_simplices_m_times(probs, m_times)
    sample_probs = probs[np.arange(n), actions.T].T

    return np.array(actions), np.array(sample_probs)

  def __str__(self):
    """Returns a string representation of a policy with parametrization."""
    return ("Softmax (linear potential): %s max`d by GA (T=%s, eta=%s)" %
            (self.obj_type, self.steps, self.step_size))
