#   Copyright (c) 2021 PaddlePaddle Authors. 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.

import parl
import paddle
import paddle.nn.functional as F
from paddle.distribution import Categorical
import numpy as np

__all__ = ['A2C']


class A2C(parl.Algorithm):
    def __init__(self, model, vf_loss_coeff=None):
        """ A2C algorithm

        Args:
            model (parl.Model): forward network of policy and value
            vf_loss_coeff (float): coefficient of the value function loss
        """

        self.model = model
        assert isinstance(vf_loss_coeff, (int, float))
        self.vf_loss_coeff = vf_loss_coeff
        clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=40.0)
        self.optimizer = paddle.optimizer.Adam(
            learning_rate=0.001,
            parameters=self.model.parameters(),
            grad_clip=clip)

    def learn(self, obs, actions, advantages, target_values, learning_rate,
              entropy_coeff):
        """
        Args:
            obs: An float32 tensor of shape ([B] + observation_space).
                 E.g. [B, C, H, W] in atari.
            actions: An int64 tensor of shape [B].
            advantages: A float32 tensor of shape [B].
            target_values: A float32 tensor of shape [B].
            learning_rate: float scalar of leanring rate.
            entropy_coeff: float scalar of entropy coefficient.
        """
        # shape: [B, act_dim]
        logits = self.model.policy(obs)
        act_dim = logits.shape[-1]
        actions_onehot = F.one_hot(actions, act_dim)
        # [B, act_dim] --> [B]
        actions_log_probs = paddle.sum(
            F.log_softmax(logits) * actions_onehot, axis=-1)
        # The policy gradient loss
        pi_loss = -1.0 * paddle.sum(actions_log_probs * advantages)

        # The value function loss
        values = self.model.value(obs)
        delta = values - target_values
        vf_loss = 0.5 * paddle.sum(paddle.square(delta))

        # The entropy loss (We want to maximize entropy, so entropy_ceoff < 0)
        # Using the Categorical just for calculating the entropy.
        # See  https://github.com/PaddlePaddle/Paddle/blob/release/2.0/python/paddle/distribution.py for detail
        policy_distribution = Categorical(logits)
        policy_entropy = policy_distribution.entropy()
        entropy = paddle.sum(policy_entropy)

        total_loss = (
            pi_loss + vf_loss * self.vf_loss_coeff + entropy * entropy_coeff)
        self.optimizer.set_lr(learning_rate)
        total_loss.backward()
        self.optimizer.step()
        self.optimizer.clear_grad()
        return total_loss, pi_loss, vf_loss, entropy

    def prob_and_value(self, obs):
        """
        Args:
            obs: An float32 tensor of shape ([B] + observation_space).
                 E.g. [B, C, H, W] in atari.
        """
        logits, values = self.model.policy_and_value(obs)
        probs = F.softmax(logits)

        return probs, values

    def predict(self, obs):
        """
        Args:
            obs: An float32 tensor of shape ([B] + observation_space).
                 E.g. [B, C, H, W] in atari.
        """
        logits = self.model.policy(obs)
        probs = F.softmax(logits)
        predict_actions = paddle.argmax(probs, 1)
        return predict_actions

    def value(self, obs):
        """
        Args:
            obs: An float32 tensor of shape ([B] + observation_space).
                 E.g. [B, C, H, W] in atari.
        """
        values = self.model.value(obs)
        return values
