# Copyright (c) 2019 Horizon Robotics. 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.
"""Various data structures."""
import collections

import numpy as np
import tensorflow as tf

from tf_agents.trajectories.policy_step import PolicyStep
from tf_agents.trajectories.time_step import StepType, TimeStep


def namedtuple(typename, field_names, default_value=None, default_values=()):
    """namedtuple with default value.

    Args:
        typename (str): type name of this namedtuple
        field_names (list[str]): name of each field
        default_value (Any): the default value for all fields
        default_values (list|dict): default value for each field
    Returns:
        the type for the namedtuple
    """
    T = collections.namedtuple(typename, field_names)
    T.__new__.__defaults__ = (default_value, ) * len(T._fields)
    if isinstance(default_values, collections.Mapping):
        prototype = T(**default_values)
    else:
        prototype = T(*default_values)
    T.__new__.__defaults__ = tuple(prototype)
    return T


class ActionTimeStep(
        namedtuple(
            'ActionTimeStep',
            [
                'step_type',  # one of (StepType.FIRST, StepType.MID, StepType.LAST)
                'reward',  # reward from the environment after executing `prev_action`
                'discount',  # discount value
                'observation',  # (nested) Tensor for observation
                'prev_action',  # action generated by the algorithm from the previous time step
                'env_id',  # the ID of the environment from whihc this time_step is from
            ])):
    """TimeStep with action."""

    def is_first(self):
        if tf.is_tensor(self.step_type):
            return tf.equal(self.step_type, StepType.FIRST)
        return np.equal(self.step_type, StepType.FIRST)

    def is_mid(self):
        if tf.is_tensor(self.step_type):
            return tf.equal(self.step_type, StepType.MID)
        return np.equal(self.step_type, StepType.MID)

    def is_last(self):
        if tf.is_tensor(self.step_type):
            return tf.equal(self.step_type, StepType.LAST)
        return np.equal(self.step_type, StepType.LAST)


def make_action_time_step(time_step: TimeStep, prev_action, first_env_id=0):
    """Make ActionTimeStep from TimeStep.

    Args:
        time_step (TimeStep):
        prev_action (nested Tensor):
        first_env_id (int): the environment ID for the first sample in this
            batch. The environment IDs for the other samples are incremented
            from this.
    Returns:
        ActionTimeStep
    """
    return ActionTimeStep(
        step_type=time_step.step_type,
        reward=time_step.reward,
        discount=time_step.discount,
        observation=time_step.observation,
        prev_action=prev_action,
        env_id=first_env_id + tf.range(time_step.step_type.shape[0]))


TrainingInfo = namedtuple(
    "TrainingInfo",
    [
        "action",
        "step_type",
        "reward",
        "discount",

        # For on-policy training, it's the PolicyStep.info from rollout
        # For off-policy training, it's the PolicyStep.info from train_step
        "info",

        # Only used for off-policy training. It's the PolicyStep.info from rollout
        "rollout_info",
        "env_id"
    ],
    default_value=())

Experience = namedtuple(
    "Experience",
    [
        'step_type',
        'reward',
        'discount',
        'observation',
        'prev_action',
        'env_id',
        'action',
        'rollout_info',  # PolicyStep.info from rollout()
        'state'  # state passed to rollout() to generate `action`
    ])


def make_experience(time_step: ActionTimeStep, policy_step: PolicyStep, state):
    """Make an instance of Experience from ActionTimeStep and PolicyStep.

    Args:
        time_step (ActionTimeStep): time step from the environment
        policy_step (PolicyStep): policy step returned from rollout()
        state (nested Tensor): state used for calling rollout() to get the
            `policy_step`
    Returns:
        Experience
    """
    return Experience(
        step_type=time_step.step_type,
        reward=time_step.reward,
        discount=time_step.discount,
        observation=time_step.observation,
        prev_action=time_step.prev_action,
        env_id=time_step.env_id,
        action=policy_step.action,
        rollout_info=policy_step.info,
        state=state)


LossInfo = namedtuple(
    "LossInfo",
    [
        "loss",  # batch loss shape should be (T, B) or (B,)
        "scalar_loss",  # shape is ()
        "extra"  # nested batch and/or scalar losses, for summary only
    ],
    default_value=())
