# coding=utf-8
# Copyright 2021 Google Health Research.
#
# 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.

"""Defines Sub-Network Routing modules and utilities."""
from typing import Callable, Dict, List, Mapping, Optional, Text, Tuple

from absl import logging
from ehr_prediction_modeling import types
from ehr_prediction_modeling.models import model_utils
from ehr_prediction_modeling.models.nets import net_utils
from ehr_prediction_modeling.utils import distributions
import sonnet as snt
import tensorflow.compat.v1 as tf
import tf_slim as slim

from ehr_prediction_modeling import configdict


class SNRConnection(snt.AbstractModule):
  """Defining the learnable connections."""

  def __init__(self,
               is_training: bool,
               conn_config: configdict.ConfigDict,
               batch_size: Optional[int] = None,
               name: str = "routing"):
    super().__init__(name=name)

    gamma_over_zeta = -conn_config.gamma / conn_config.zeta
    if gamma_over_zeta <= 0:
      raise ValueError(
          f"-Gamma/zeta needs to be greater than 0. Found {gamma_over_zeta}")

    self._is_training = is_training
    self._batch_size = batch_size
    self._subnetwork_conn_type = conn_config.subnetwork_to_subnetwork_conn_type
    self._reg_factor_weight = conn_config.subnetwork_conn_l_reg_factor_weight
    self._var_name = name
    with self._enter_variable_scope():
      self._log_alpha = tf.get_variable(
          self._var_name, shape=[], initializer=tf.constant_initializer(-0.5))
      self._distribution = distributions.HardConcrete(
          self._log_alpha,
          temperature=conn_config.beta,
          lower=conn_config.gamma,
          higher=conn_config.zeta,
          validate_args=True,
      )

  def get_regularization_loss(self) -> tf.Tensor:
    """Returns the regularization loss corresponding to an L0 penalty term."""
    if self._reg_factor_weight <= 0:
      return tf.constant(0.)

    l_reg_factor_weight = tf.constant(self._reg_factor_weight)
    return l_reg_factor_weight * self._distribution.l0_norm()

  def _build(self, inp: tf.Tensor) -> tf.Tensor:
    """Creates a learnable variable for constructing connection of sub_networks.

    Edge will be generated by a hard sigmoid function bounded by gamma and zeta,
    with temperature set to beta.

    Args:
      inp: The input to connect.

    Returns:
      The input scaled by the connection.
    """
    if self._subnetwork_conn_type == types.SubNettoSubNetConnType.SCALAR_WEIGHT:
      return self._log_alpha * inp

    if self._is_training:
      batch_size = () if self._batch_size is None else (self._batch_size, 1)
      return self._distribution.sample(batch_size) * inp
    else:
      return self._distribution.hard_sigmoid() * inp


def all_tensors_equal_final_dim(inputs: tf.Tensor) -> bool:
  """Checks whether all the tensors in a list are of same shape."""
  return all(x.get_shape().as_list()[-1] == inputs[0].get_shape().as_list()[-1]
             for x in inputs)


def combine_inputs(inputs: List[tf.Tensor],
                   input_combination_method: str) -> Optional[tf.Tensor]:
  """Computes the inputs of a subnetwork."""
  if input_combination_method == (types.SNRInputCombinationType.CONCATENATE):
    combined_input = tf.concat(inputs, axis=-1)
  elif input_combination_method == (types.SNRInputCombinationType.SUM_ALL):
    if not all_tensors_equal_final_dim(inputs):
      raise ValueError("Input tensors are of unequal shape. Hence "
                       "types.SNRInputCombinationType.SUM_ALL can't be "
                       "applied. Debug or use method of combining the inputs "
                       "listed in types.SNRInputCombinationType.")
    combined_input = tf.add_n(inputs)
  else:
    raise ValueError("Unknown subnetwork input combination method. Use one of "
                     "types.SNRInputCombinationType.")
  return combined_input


class SNRTaskLayer(snt.AbstractModule):
  """Defines the layer between model output and logits using SNR connections."""

  routing_connection_name = "task_routing"

  def __init__(self,
               name: Text,
               output_sizes: List[int],
               initializers: Mapping[Text, Callable[[tf.Tensor], tf.Tensor]],
               regularizer: Optional[Callable[[List[tf.Variable]], tf.Tensor]],
               activation: Callable[[tf.Tensor], tf.Tensor],
               is_training: bool = False,
               snr_config: Optional[configdict.ConfigDict] = None) -> None:
    super(SNRTaskLayer, self).__init__(name=f"task_layer_{name}")
    self._is_training = is_training
    self._snr_config = snr_config
    self._regularizer = regularizer
    self._task_name = name
    task_layer_kwargs = {
        "output_sizes": output_sizes,
        "initializers": initializers,
        "activation": activation,
        "activate_final": False,
        "name": f"{self._task_name}_logistic"
    }
    self._mlp = snt.nets.MLP(**task_layer_kwargs)
    self._snr_connections = []
    logging.info("Created SNR task layer %s with is_training %s and args %s",
                 name, is_training, task_layer_kwargs)

  def _build(self, model_output: tf.Tensor) -> tf.Tensor:
    # When using SNRNN the model_output contains a list of tensors, each
    # corresponding to the output of one cell in the model. Those outputs will
    # be combined with SNRConnections, boolean trainable variables, before being
    # passed through the fully connected layer.
    assert isinstance(model_output, list)
    assert self._snr_config

    inputs = []
    for i, output in enumerate(model_output):
      with tf.variable_scope(f"task_layer_{self._task_name}_input_{i}"):
        routing = SNRConnection(
            is_training=self._is_training,
            name=self.routing_connection_name,
            conn_config=self._snr_config)
        self._snr_connections.append(routing)
        inputs.append(routing(output))
    task_input = combine_inputs(inputs,
                                self._snr_config.input_combination_method)

    return snt.BatchApply(self._mlp)(task_input)

  def get_snr_connections_loss(self) -> tf.Tensor:
    """Gets the regularization loss for all the routing connections."""
    if not self._snr_config:
      return tf.constant(0.)
    return sum(
        [conn.get_regularization_loss() for conn in self._snr_connections])

  def get_layer_regularization_loss(self) -> tf.Tensor:
    """Gets the regularization loss for all task layer variables."""
    snr_regularization_penalty = self.get_snr_connections_loss()

    if not self._regularizer:
      return tf.constant(0.)
    logistic_penalty = slim.layers.regularizers.apply_regularization(
        self._regularizer, self._mlp.get_all_variables())

    return snr_regularization_penalty + logistic_penalty


class SNREncoderLayer(snt.AbstractModule):
  """Builds a sub-network routing layer for the SNREncoder."""

  def __init__(self,
               subnetwork_coordinates: Tuple[int, int],
               batch_size: int,
               activation_function: Callable[[tf.Tensor], tf.Tensor],
               snr_config: configdict.ConfigDict,
               nhidden: int,
               tasks: List[Text],
               initializers: Optional[Dict[Text, tf.Tensor]] = None,
               dropout_prob: float = 0.,
               use_bias: bool = True,
               is_training: bool = True,
               is_final_layer: bool = False,
               activate_final: bool = False) -> None:
    super(SNREncoderLayer, self).__init__(
        name=f"enc_layer{subnetwork_coordinates[0]}{subnetwork_coordinates[1]}")
    self._subnetwork_coordinates = subnetwork_coordinates
    self._batch_size = batch_size
    self._activation_function = activation_function
    self._snr_config = snr_config
    self._nhidden = nhidden
    self._tasks = sorted(tasks)
    self._initializers = initializers
    self._dropout_prob = dropout_prob
    self._use_bias = use_bias
    self._is_training = is_training
    self._is_final_layer = is_final_layer
    self._activate_final = activate_final
    self._snr_connections = []
    self._linear_layers = []

  def _get_routed_input(self, layer: int, inp: tf.Tensor, is_training: bool,
                        snr_config: configdict.ConfigDict,
                        batch_size: int) -> tf.Tensor:
    # If task specific routing is not used then there will be only one
    # connection created for the given input.
    # Alternatively if the task list isn't passed in the config then task
    # specific routing is considered as disabled.
    if (not self._snr_config.get("use_task_specific_routing", False) or
        not self._tasks):
      routing = SNRConnection(
          is_training=is_training,
          conn_config=snr_config,
          batch_size=batch_size)
      self._snr_connections.append(routing)
      return inp if layer == 0 else routing(inp)

    subnet_activations = []
    for task_name in self._tasks:
      with tf.variable_scope(f"task_{task_name}"):
        routing = SNRConnection(
            is_training=is_training,
            conn_config=snr_config,
            batch_size=batch_size)
        self._snr_connections.append(routing)
        # We don't use Boolean connections between the output
        # of the SparseLookupEmbedding layer and the first
        # layer of sub-networks.
        subnet_act_for_curr_input = inp if layer == 0 else routing(inp)
        subnet_activations.append(subnet_act_for_curr_input)
    return combine_inputs(subnet_activations,
                          snr_config.input_combination_method)

  def get_snr_connections_loss(self) -> tf.Tensor:
    """Gets the regularization loss for all the routing connections."""
    if not self._snr_config:
      return tf.constant(0.)
    return sum(
        [conn.get_regularization_loss() for conn in self._snr_connections])

  def get_weights_penalty(self) -> tf.Tensor:
    snr_encoder_regularization_penalty = tf.constant(0.)
    if self._snr_config.subnetwork_weight_l_reg_factor_weight > 0:
      snr_encoder_weights = [
          linear_layer.w for linear_layer in self._linear_layers
      ]
      l_reg_factor_weight = tf.constant(
          self._snr_config.subnetwork_weight_l_reg_factor_weight)
      regularizer = model_utils.get_regularizer(
          self._snr_config.subnetwork_weight_l_reg, l_reg_factor_weight)
      snr_subnetwork_weight_penalty = (
          slim.layers.regularizers.apply_regularization(regularizer,
                                                        snr_encoder_weights))
      snr_encoder_regularization_penalty += snr_subnetwork_weight_penalty

    return snr_encoder_regularization_penalty

  def _build(self, inputs: List[tf.Tensor]) -> tf.Tensor:
    layer_index, unit_index = self._subnetwork_coordinates

    if not self._linear_layers:
      for k in range(len(inputs)):
        linear_layer = snt.Linear(
            output_size=self._nhidden,
            use_bias=self._use_bias,
            name=f"layer_{layer_index}_subnetwork_{unit_index}_input_{k}")
        self._linear_layers.append(linear_layer)

    subnetwork_input = None
    for k, inp in enumerate(inputs):
      with tf.variable_scope(f"input_{k}"):
        linear_inp = self._linear_layers[k](inp)
        h = self._get_routed_input(
            layer=layer_index,
            inp=linear_inp,
            is_training=self._is_training,
            snr_config=self._snr_config,
            batch_size=self._batch_size)
        if not self._is_final_layer and self._dropout_prob > 0:
          h = slim.layers.dropout(
              h,
              keep_prob=(1 - self._dropout_prob),
              initializers=self._initializers,
              is_training=self._is_training)
        if subnetwork_input is None:
          if (self._is_final_layer and self._activate_final) or (
              self._snr_config.activation_before_aggregation):
            subnetwork_input = self._activation_function(h)
          else:
            subnetwork_input = h
        else:
          if (self._is_final_layer and self._activate_final) or (
              self._snr_config.activation_before_aggregation):
            subnetwork_input += self._activation_function(h)
          else:
            subnetwork_input += h

    if self._is_final_layer:
      return subnetwork_input

    if not self._snr_config.activation_before_aggregation:
      subnetwork_input = self._activation_function(subnetwork_input)

    if self._snr_config.snr_block_conn_type == types.SNRBlockConnType.NONE:
      return subnetwork_input

    y = self._activation_function(
        snt.Linear(output_size=self._nhidden,
                   use_bias=self._use_bias)(subnetwork_input))
    if self._snr_config.snr_block_conn_type == types.SNRBlockConnType.FC:
      return y
    elif (self._snr_config.snr_block_conn_type ==
          types.SNRBlockConnType.RESIDUAL):
      return y + subnetwork_input
    elif (
        self._snr_config.snr_block_conn_type == types.SNRBlockConnType.HIGHWAY):
      # If highway, then mix data and transform using convex combination
      # Link to paper: https://arxiv.org/abs/1505.00387
      transform_gate = tf.sigmoid(
          snt.Linear(
              output_size=self._nhidden,
              initializers=net_utils.INITIALIZER_DICT)(subnetwork_input))
      return tf.multiply(y, transform_gate) + tf.multiply(
          subnetwork_input, 1. - transform_gate)
    else:
      raise ValueError(
          "Wrong SNREncoder unit type. Use one of types.SNRBlockConnType.")
