# This file is part of tf-mdp.

# tf-mdp is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# tf-mdp is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with tf-mdp. If not, see <http://www.gnu.org/licenses/>.


import rddl2tf

from tfmdp.policy.drp import DeepReactivePolicy

from tfmdp.model.cell.reparameterization_cell import ReparameterizationCell
from tfmdp.model.sequential.mrm import MarkovRecurrentModel, Trajectory

from tfmdp.model import utils

import tensorflow as tf
from typing import Dict, Sequence, Optional, Tuple


class ReparameterizationSampling(MarkovRecurrentModel):
    '''ReparameterizationSampling builds the symbolic computation graph
    for approximating the finite-horizon expected return of a policy
    for a given initial state by means of Monte-Carlo approximation.
    The trajectories are generated by stochastic reparameterization
    defined over noise vectors given as inputs.

    Args:
        compiler (:obj:`rddl2tf.compilers.Compiler`): RDDL2TensorFlow compiler.
        config (Dict): The recurrent model configuration parameters.
    '''

    def __init__(self,
                 compiler: rddl2tf.compilers.Compiler,
                 config: Dict) -> None:
        super(ReparameterizationSampling, self).__init__(compiler, config)

    def build(self, policy: DeepReactivePolicy) -> None:
        '''Builds a basic Markov cell ops by embedding the `policy` in the transition sampling.

        Args:
            policy (:obj:`tfmdp.policy.drp.DeepReactivePolicy`): A deep reactive policy.
        '''
        self.reparameterization_map = self.compiler.get_cpfs_reparameterization()
        self.cell = ReparameterizationCell(self.compiler, policy, config={})

    def __call__(self,
                 initial_state: Sequence[tf.Tensor],
                 horizon: int) -> Tuple[Trajectory, Sequence[tf.Tensor], tf.Tensor]:
        '''Samples a batch state-action-reward trajectory with given
        `initial_state` and `horizon`, and returns the corresponding total reward.

        Args:
            initial_state (Sequence[tf.Tensor]): The initial state tensors.
            horizon (int): The number of timesteps in each sampled trajectory.

        Returns:
            Tuple[Trajectory, Sequence[tf.Tensor], tf.Tensor]: A triple of (namedtuple, tensors, tensor)
            representing the trajectory, final state, and total reward.
        '''
        batch_size = int(initial_state[0].shape[0])

        with self.graph.as_default():

            with tf.compat.v1.variable_scope('reparameterization'):
                self.noise_map = utils.get_noise_variables(self.reparameterization_map, batch_size, horizon)
                self.noise, encoding = utils.encode_noise_as_inputs(self.noise_map)
                self.cell.config['encoding'] = encoding

            with tf.compat.v1.name_scope('inputs'):
                self.timesteps = self.timesteps(horizon, batch_size)
                self.inputs = tf.concat([self.timesteps, self.noise], axis=2)

            with tf.compat.v1.name_scope('trajectory'):
                outputs, final_state = tf.nn.dynamic_rnn(self.cell,
                                                         self.inputs,
                                                         initial_state=initial_state,
                                                         dtype=tf.float32)

                states = tuple(fluent[0] for fluent in outputs[0])
                actions = tuple(fluent[0] for fluent in outputs[1])
                interms = tuple(fluent[0] for fluent in outputs[2])
                rewards = outputs[3]
                trajectory = Trajectory(states, actions, interms, rewards)

            with tf.compat.v1.name_scope('total_reward'):
                total_reward = tf.reduce_sum(tf.squeeze(trajectory.rewards), axis=1)

        return (trajectory, final_state, total_reward)

    def summary(self) -> None:
        '''Prints a string summary of the recurrent model.'''
        raise NotImplementedError

    @property
    def trainable_variables(self) -> Sequence[tf.Variable]:
        '''Returns the list of trainable variables.'''
        with self.graph.as_default():
            noise_variables = tf.compat.v1.trainable_variables(r'.*reparameterization')
            return noise_variables
