# Copyright 2020 Google LLC.
#
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v2 as tf

from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union

from .dataset import Dataset, EnvStep, OffpolicyDataset, OnpolicyDataset, StepType
from . import common as common_lib


def _default_by_steps_reward_fn(env_step):
  return env_step.reward


def _default_by_episodes_reward_fn(env_step, valid_steps, gamma):
  mask = (
      (1 - tf.cast(env_step.is_last(), tf.float32)) *
      tf.cast(valid_steps, tf.float32))
  discount = tf.pow(tf.cast(gamma, tf.float32),
                    tf.cast(env_step.step_num, tf.float32))
  return tf.reduce_sum(mask * discount * env_step.reward, -1)


def _default_by_steps_weight_fn(env_step, gamma):
  mask = 1 - tf.cast(env_step.is_last(), tf.float32)
  discount = tf.pow(tf.cast(gamma, tf.float32),
                    tf.cast(env_step.step_num, tf.float32))
  return mask * discount


def _default_by_episodes_weight_fn(env_step, valid_steps):
  return tf.ones([tf.shape(valid_steps)[0]], dtype=tf.float32)


def get_minibatch_average(dataset: Dataset,
                          batch_size: int,
                          num_batches: int = 1,
                          by_steps: bool = True,
                          truncate_episode_at: Optional[int] = None,
                          reward_fn: Callable = None,
                          weight_fn: Callable = None,
                          gamma: Union[float, tf.Tensor] = 1.0) -> Union[
                                  float, tf.Tensor]:
  """Computes average reward via randomly sampled mini-batches.

    Samples steps or episodes from the dataset and computes average reward.

    Args:
      dataset: The dataset to sample experience from.
      batch_size: The number of episodes to sample per batch.
      num_batches: The number of batches to use for estimation.
      by_steps: Whether to sample batches of steps (default) or episodes.
      truncate_episode_at: If sampling by episodes, where to truncate episodes
        from the environment, if at all.
      reward_fn: A function that takes in an EnvStep and returns the reward for
        that step. If not specified, defaults to just EnvStep.reward. When
        sampling by episode, valid_steps is also passed into reward_fn.
      weight_fn: A function that takes in an EnvStep and returns a weight for
        that step. If not specified, defaults to gamma ** step_num. When
        sampling by episode, valid_steps is also passed into reward_fn.
      gamma: The discount factor to use for the default reward/weight functions.

    Returns:
      An estimate of the average reward.
    """
  if reward_fn is None:
    if by_steps:
      reward_fn = _default_by_steps_reward_fn
    else:
      reward_fn = lambda *args: _default_by_episodes_reward_fn(
          *args, gamma=gamma)

  if weight_fn is None:
    if by_steps:
      weight_fn = lambda *args: _default_by_steps_weight_fn(*args, gamma=gamma)
    else:
      weight_fn = _default_by_episodes_weight_fn

  total_reward = 0.
  total_weight = 0.
  for _ in range(num_batches):
    if by_steps:
      if isinstance(dataset, OnpolicyDataset):
        steps = dataset.get_step(num_steps=batch_size)
      else:
        steps = dataset.get_step(batch_size)
      rewards = reward_fn(steps)
      weights = weight_fn(steps)
    else:
      episodes, valid_steps = dataset.get_episode(
          batch_size, truncate_episode_at=truncate_episode_at)
      rewards = reward_fn(episodes, valid_steps)
      weights = weight_fn(episodes, valid_steps)

    rewards = common_lib.reverse_broadcast(rewards, weights)
    weights = common_lib.reverse_broadcast(weights, rewards)
    total_reward += tf.reduce_sum(rewards * weights, axis=0)
    total_weight += tf.reduce_sum(weights, axis=0)

  return total_reward / total_weight


def get_fullbatch_average(dataset: OffpolicyDataset,
                          limit: Optional[int] = None,
                          by_steps: bool = True,
                          truncate_episode_at: Optional[int] = None,
                          reward_fn: Callable = None,
                          weight_fn: Callable = None,
                          gamma: Union[float, tf.Tensor] = 1.0) -> Union[
                                  float, tf.Tensor]:
  """Computes average reward over full dataset.

    Args:
      dataset: The dataset to sample experience from.
      limit: If specified, the maximum number of steps/episodes to take from the
        dataset.
      by_steps: Whether to sample batches of steps (default) or episodes.
      truncate_episode_at: If sampling by episodes, where to truncate episodes
        from the environment, if at all.
      reward_fn: A function that takes in an EnvStep and returns the reward for
        that step. If not specified, defaults to just EnvStep.reward. When
        sampling by episode, valid_steps is also passed into reward_fn.
      weight_fn: A function that takes in an EnvStep and returns a weight for
        that step. If not specified, defaults to gamma ** step_num. When
        sampling by episode, valid_steps is also passed into reward_fn.
      gamma: The discount factor to use for the default reward/weight functions.

    Returns:
      An estimate of the average reward.
    """
  if reward_fn is None:
    if by_steps:
      reward_fn = _default_by_steps_reward_fn
    else:
      reward_fn = lambda *args: _default_by_episodes_reward_fn(
          *args, gamma=gamma)

  if weight_fn is None:
    if by_steps:
      weight_fn = lambda *args: _default_by_steps_weight_fn(*args, gamma=gamma)
    else:
      weight_fn = _default_by_episodes_weight_fn

  if by_steps:
    steps = dataset.get_all_steps(limit=limit)
    rewards = reward_fn(steps)
    weights = weight_fn(steps)
  else:
    episodes, valid_steps = dataset.get_all_episodes(
        truncate_episode_at=truncate_episode_at, limit=limit)
    rewards = reward_fn(episodes, valid_steps)
    weights = weight_fn(episodes, valid_steps)

  rewards = common_lib.reverse_broadcast(rewards, weights)
  weights = common_lib.reverse_broadcast(weights, rewards)
  if tf.rank(weights) < 2:
    return (tf.reduce_sum(rewards * weights, axis=0) /
            tf.reduce_sum(weights, axis=0))
  return (tf.linalg.matmul(weights, rewards) /
          tf.reduce_sum(tf.math.reduce_mean(weights, axis=0)))
