# Copyright 2017 reinforce.io. 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.
# ==============================================================================

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

import tensorflow as tf

from tensorforce.core.baselines import Baseline, AggregatedBaseline
from tensorforce.core.optimizers import Optimizer
from tensorforce.models import DistributionModel


class PGModel(DistributionModel):
    """
    Base class for policy gradient models. It optionally defines a baseline
    and handles its optimization. It implements the `tf_loss_per_instance` function, but requires
    subclasses to implement `tf_pg_loss_per_instance`.
    """

    COMPONENT_BASELINE = "baseline"

    def __init__(
        self,
        states,
        actions,
        scope,
        device,
        saver,
        summarizer,
        execution,
        batching_capacity,
        variable_noise,
        states_preprocessing,
        actions_exploration,
        reward_preprocessing,
        update_mode,
        memory,
        optimizer,
        discount,
        network,
        distributions,
        entropy_regularization,
        baseline_mode,
        baseline,
        baseline_optimizer,
        gae_lambda
    ):
        # Baseline mode
        assert baseline_mode is None or baseline_mode in ('states', 'network')
        self.baseline_mode = baseline_mode

        self.baseline_spec = baseline
        self.baseline_optimizer_spec = baseline_optimizer

        # Generalized advantage function
        assert gae_lambda is None or (0.0 <= gae_lambda <= 1.0 and self.baseline_mode is not None)
        self.gae_lambda = gae_lambda

        self.baseline = None
        self.baseline_optimizer = None
        self.fn_reward_estimation = None

        super(PGModel, self).__init__(
            states=states,
            actions=actions,
            scope=scope,
            device=device,
            saver=saver,
            summarizer=summarizer,
            execution=execution,
            batching_capacity=batching_capacity,
            variable_noise=variable_noise,
            states_preprocessing=states_preprocessing,
            actions_exploration=actions_exploration,
            reward_preprocessing=reward_preprocessing,
            update_mode=update_mode,
            memory=memory,
            optimizer=optimizer,
            discount=discount,
            network=network,
            distributions=distributions,
            entropy_regularization=entropy_regularization,
            requires_deterministic=False
        )

    def as_local_model(self):
        super(PGModel, self).as_local_model()
        if self.baseline_optimizer_spec is not None:
            self.baseline_optimizer_spec = dict(
                type='global_optimizer',
                optimizer=self.baseline_optimizer_spec
            )

    def initialize(self, custom_getter):
        super(PGModel, self).initialize(custom_getter)

        # Baseline
        if self.baseline_spec is None:
            assert self.baseline_mode is None

        elif all(name in self.states_spec for name in self.baseline_spec):
            # Implies AggregatedBaseline.
            assert self.baseline_mode == 'states'
            self.baseline = AggregatedBaseline(baselines=self.baseline_spec)

        else:
            assert self.baseline_mode is not None
            self.baseline = Baseline.from_spec(
                spec=self.baseline_spec,
                kwargs=dict(
                    summary_labels=self.summary_labels
                )
            )

        # Baseline optimizer
        if self.baseline_optimizer_spec is not None:
            assert self.baseline_mode is not None
            self.baseline_optimizer = Optimizer.from_spec(spec=self.baseline_optimizer_spec)

        # TODO: Baseline internal states !!! (see target_network q_model)

        # Reward estimation
        self.fn_reward_estimation = tf.make_template(
            name_='reward-estimation',
            func_=self.tf_reward_estimation,
            custom_getter_=custom_getter
        )
        # Baseline loss
        self.fn_baseline_loss = tf.make_template(
            name_='baseline-loss',
            func_=self.tf_baseline_loss,
            custom_getter_=custom_getter
        )

    def tf_reward_estimation(self, states, internals, terminal, reward, update):
        if self.baseline_mode is None:
            return self.fn_discounted_cumulative_reward(terminal=terminal, reward=reward, discount=self.discount)

        else:
            if self.baseline_mode == 'states':
                state_value = self.baseline.predict(
                    states=states,
                    internals=internals,
                    update=update
                )

            elif self.baseline_mode == 'network':
                embedding = self.network.apply(
                    x=states,
                    internals=internals,
                    update=update
                )
                state_value = self.baseline.predict(
                    states=tf.stop_gradient(input=embedding),
                    internals=internals,
                    update=update
                )

            if self.gae_lambda is None:
                reward = self.fn_discounted_cumulative_reward(
                    terminal=terminal,
                    reward=reward,
                    discount=self.discount
                )
                advantage = reward - state_value

            else:
                next_state_value = tf.concat(values=(state_value[1:], (0.0,)), axis=0)
                zeros = tf.zeros_like(tensor=next_state_value)
                next_state_value = tf.where(condition=terminal, x=zeros, y=next_state_value)
                td_residual = reward + self.discount * next_state_value - state_value
                gae_discount = self.discount * self.gae_lambda
                advantage = self.fn_discounted_cumulative_reward(
                    terminal=terminal,
                    reward=td_residual,
                    discount=gae_discount
                )

            # Normalize advantage.
            # mean, variance = tf.nn.moments(advantage, axes=[0], keep_dims=True)
            # advantage = (advantage - mean) / tf.sqrt(x=tf.maximum(x=variance, y=util.epsilon))

            return advantage

    def tf_regularization_losses(self, states, internals, update):
        losses = super(PGModel, self).tf_regularization_losses(
            states=states,
            internals=internals,
            update=update
        )

        if self.baseline_mode is not None and self.baseline_optimizer is None:
            baseline_regularization_loss = self.baseline.regularization_loss()
            if baseline_regularization_loss is not None:
                losses['baseline'] = baseline_regularization_loss

        return losses

    def tf_baseline_loss(self, states, internals, reward, update, reference=None):
        """
        Creates the TensorFlow operations for calculating the baseline loss of a batch.

        Args:
            states: Dict of state tensors.
            internals: List of prior internal state tensors.
            reward: Reward tensor.
            update: Boolean tensor indicating whether this call happens during an update.
            reference: Optional reference tensor(s), in case of a comparative loss.

        Returns:
            Loss tensor.
        """
        if self.baseline_mode == 'states':
            loss = self.baseline.loss(
                states=states,
                internals=internals,
                reward=reward,
                update=update,
                reference=reference
            )

        elif self.baseline_mode == 'network':
            loss = self.baseline.loss(
                states=self.network.apply(x=states, internals=internals, update=update),
                internals=internals,
                reward=reward,
                update=update,
                reference=reference
            )

        regularization_loss = self.baseline.regularization_loss()
        if regularization_loss is not None:
            loss += regularization_loss

        return loss

    def baseline_optimizer_arguments(self, states, internals, reward):
        """
        Returns the baseline optimizer arguments including the time, the list of variables to  
        optimize, and various functions which the optimizer might require to perform an update  
        step.

        Args:
            states: Dict of state tensors.
            internals: List of prior internal state tensors.
            reward: Reward tensor.

        Returns:
            Baseline optimizer arguments as dict.
        """
        arguments = dict(
            time=self.global_timestep,
            variables=self.baseline.get_variables(),
            arguments=dict(
                states=states,
                internals=internals,
                reward=reward,
                update=tf.constant(value=True),
            ),
            fn_reference=self.baseline.reference,
            fn_loss=self.fn_baseline_loss,
            # source_variables=self.network.get_variables()
        )
        if self.global_model is not None:
            arguments['global_variables'] = self.global_model.baseline.get_variables()
        return arguments

    def tf_optimization(self, states, internals, actions, terminal, reward, next_states=None, next_internals=None):
        assert next_states is None and next_internals is None  # temporary

        estimated_reward = self.fn_reward_estimation(
            states=states,
            internals=internals,
            terminal=terminal,
            reward=reward,
            update=tf.constant(value=True)
        )
        if self.baseline_optimizer is not None:
            estimated_reward = tf.stop_gradient(input=estimated_reward)

        optimization = super(PGModel, self).tf_optimization(
            states=states,
            internals=internals,
            actions=actions,
            terminal=terminal,
            reward=estimated_reward,
            next_states=next_states,
            next_internals=next_internals
        )

        if self.baseline_optimizer is not None:
            cumulative_reward = self.fn_discounted_cumulative_reward(terminal=terminal, reward=reward, discount=self.discount)

            arguments = self.baseline_optimizer_arguments(
                states=states,
                internals=internals,
                reward=cumulative_reward,
            )
            baseline_optimization = self.baseline_optimizer.minimize(**arguments)

            optimization = tf.group(optimization, baseline_optimization)

        return optimization

    def get_variables(self, include_submodules=False, include_nontrainable=False):
        model_variables = super(PGModel, self).get_variables(
            include_submodules=include_submodules,
            include_nontrainable=include_nontrainable
        )

        if self.baseline_mode is not None and (include_submodules or self.baseline_optimizer is None):
            baseline_variables = self.baseline.get_variables(include_nontrainable=include_nontrainable)
            model_variables += baseline_variables

            if include_nontrainable and self.baseline_optimizer is not None:
                baseline_optimizer_variables = self.baseline_optimizer.get_variables()
                # For some reason, some optimizer variables are only registered in the model.
                for variable in baseline_optimizer_variables:
                    if variable in model_variables:
                        model_variables.remove(variable)
                model_variables += baseline_optimizer_variables

        return model_variables

    def get_summaries(self):
        if self.baseline_mode is None:
            return super(PGModel, self).get_summaries()
        else:
            return super(PGModel, self).get_summaries() + self.baseline.get_summaries()

    def get_components(self):
        if self.baseline is None:
            return super(PGModel, self).get_components()
        else:
            result = dict(super(PGModel, self).get_components())
            result[PGModel.COMPONENT_BASELINE] = self.baseline
            return result
