# ==============================================================================
# Copyright 2019 - Philip Paquette
#
# NOTICE:  Permission is hereby granted, free of charge, to any person obtaining
#   a copy of this software and associated documentation files (the "Software"),
#   to deal in the Software without restriction, including without limitation the
#   rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
#   sell copies of the Software, and to permit persons to whom the Software is
#   furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in all
#   copies or substantial portions of the Software.
# ==============================================================================
""" Policy model (v002_markovian_film)
    - Contains the policy model (v002_markovian_film), to evaluate the best actions given a state
"""
import logging
from diplomacy_research.models.policy.token_based import TokenBasedPolicyModel, load_args as load_parent_args
from diplomacy_research.models.state_space import get_adjacency_matrix, VOCABULARY_SIZE, TOKENS_PER_ORDER, \
    NB_SUPPLY_CENTERS, NB_POWERS

# Constants
LOGGER = logging.getLogger(__name__)

def load_args():
    """ Load possible arguments
        :return: A list of tuple (arg_type, arg_name, arg_value, arg_desc)
    """
    return load_parent_args() + [
        # Hyperparameters
        ('int', 'nb_graph_conv', 8, 'Number of Graph Conv Layer'),
        ('int', 'word_emb_size', 400, 'Word embedding size.'),
        ('int', 'power_emb_size', 64, 'Power embedding size.'),
        ('int', 'gcn_size', 80, 'Size of graph convolution outputs.'),
        ('int', 'lstm_size', 400, 'LSTM (Encoder and Decoder) size.'),
        ('int', 'attn_size', 80, 'LSTM decoder attention size.'),
    ]

class PolicyModel(TokenBasedPolicyModel):
    """ Policy Model """

    def _encode_board(self, board_state, name, reuse=None):
        """ Encodes a board state or prev orders state
            :param board_state: The board state / prev orders state to encode - (batch, NB_NODES, initial_features)
            :param name: The name to use for the encoding
            :param reuse: Whether to reuse or not the weights from another encoding operation
            :return: The encoded board state / prev_orders state
        """
        from diplomacy_research.utils.tensorflow import tf
        from diplomacy_research.models.layers.graph_convolution import film_gcn_res_block, preprocess_adjacency

        # Quick function to retrieve hparams and placeholders and function shorthands
        hps = lambda hparam_name: self.hparams[hparam_name]
        pholder = lambda placeholder_name: self.placeholders[placeholder_name]
        relu = tf.nn.relu

        # Getting film gammas and betas
        film_gammas = self.outputs['_%s_film_gammas' % name]
        film_betas = self.outputs['_%s_film_betas' % name]

        # Computing norm adjacency
        norm_adjacency = preprocess_adjacency(get_adjacency_matrix())
        norm_adjacency = tf.tile(tf.expand_dims(norm_adjacency, axis=0), [tf.shape(board_state)[0], 1, 1])

        # Building scope
        scope = tf.VariableScope(name='policy/%s' % name, reuse=reuse)
        with tf.variable_scope(scope):

            # Adding noise to break symmetry
            board_state = board_state + tf.random_normal(tf.shape(board_state), stddev=0.01)
            graph_conv = tf.layers.Dense(units=hps('gcn_size'), activation=relu)(board_state)

            # First and intermediate layers
            for layer_idx in range(hps('nb_graph_conv') - 1):
                graph_conv = film_gcn_res_block(inputs=graph_conv,                          # (b, NB_NODES, gcn_size)
                                                gamma=film_gammas[layer_idx],
                                                beta=film_betas[layer_idx],
                                                gcn_out_dim=hps('gcn_size'),
                                                norm_adjacency=norm_adjacency,
                                                is_training=pholder('is_training'),
                                                residual=True)

            # Last layer
            graph_conv = film_gcn_res_block(inputs=graph_conv,                              # (b, NB_NODES, attn_size)
                                            gamma=film_gammas[-1],
                                            beta=film_betas[-1],
                                            gcn_out_dim=hps('attn_size'),
                                            norm_adjacency=norm_adjacency,
                                            is_training=pholder('is_training'),
                                            residual=False)

        # Returning
        return graph_conv