# coding=utf-8
# Copyright 2019 The Edward2 Authors.
#
# 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.

"""Initializers.

This module extends `tf.keras.initializers` with the notion of "trainable
initializers", where initializers to weights and biases in `tf.keras.layers` may
themselves carry parameters. For example, consider a weight initializer which
returns a variational distribution: this is reified as an `ed.RandomVariable`
parameterized by `tf.Variables`.

One subtlety is how `tf.keras.constraints` are used on the parameters of
trainable initializers. Typically, Keras constraints are used with projected
gradient descent, where one performs unconstrained optimization and then applies
a projection (the constraint) after each gradient update. To stay in line with
probabilistic literature, trainable initializers apply constraints on the
`tf.Variables` themselves (i.e., a constrained parameterization) and do not
apply projections during optimization.
"""

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

import math

from edward2.tensorflow import constraints
from edward2.tensorflow import generated_random_variables
from edward2.tensorflow import regularizers
import six
import tensorflow.compat.v1 as tf1
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp


# From `tensorflow/python/ops/init_ops.py`
def _compute_fans(shape):
  """Computes the number of input and output units for a weight shape.

  Args:
    shape: Integer shape tuple or TF tensor shape.

  Returns:
    A tuple of scalars (fan_in, fan_out).
  """
  if len(shape) < 1:  # Just to avoid errors for constants.
    fan_in = fan_out = 1
  elif len(shape) == 1:
    fan_in = fan_out = shape[0]
  elif len(shape) == 2:
    fan_in = shape[0]
    fan_out = shape[1]
  else:
    # Assuming convolution kernels (2D, 3D, or more).
    # kernel shape: (..., input_depth, depth)
    receptive_field_size = 1.
    for dim in shape[:-2]:
      receptive_field_size *= dim
    fan_in = shape[-2] * receptive_field_size
    fan_out = shape[-1] * receptive_field_size
  if isinstance(fan_in, tf1.Dimension):
    fan_in = fan_in.value
  if isinstance(fan_out, tf1.Dimension):
    fan_out = fan_out.value
  return fan_in, fan_out


class ScaledNormalStdDev(tf.keras.initializers.VarianceScaling):
  """Initializer capable of adapting its scale to the shape of weights tensors.

  This initializes the standard deviation parameter of a Trainable Normal
  distribution with a scale based on the shape of the weights tensor.
  Additionally, A small amount of noise will be added to break weigh symmetry.

  With `distribution="truncated_normal" or "untruncated_normal"`, the standard
  deviation (after truncation, if used) is `stddev = sqrt(scale / n)`, where n
  is:
    - number of input units in the weight tensor, if mode = "fan_in"
    - number of output units, if mode = "fan_out"
    - average of the numbers of input and output units, if mode = "fan_avg"
  """

  def __init__(self,
               scale=1.0,
               mode='fan_in',
               distribution='untruncated_normal',
               seed=None):
    """Constructs the initializer.

    Args:
      scale: Scaling factor (positive float).
      mode: One of "fan_in", "fan_out", "fan_avg".
      distribution: Random distribution to use. One of "truncated_normal", or
        "untruncated_normal".
      seed: A Python integer. Used to create random seeds. See
        `tf.set_random_seed`
        for behavior.

    Raises:
      ValueError: In case of an invalid value for the "scale", mode" or
        "distribution" arguments.
    """
    distribution = distribution.lower()
    if distribution not in {'truncated_normal', 'untruncated_normal'}:
      raise ValueError('Invalid `distribution` argument:', distribution)
    super(ScaledNormalStdDev, self).__init__(scale=scale, mode=mode,
                                             distribution=distribution,
                                             seed=seed)

  def __call__(self, shape, dtype=None):
    if dtype is None:
      dtype = self.dtype
    scale = self.scale
    scale_shape = shape
    fan_in, fan_out = _compute_fans(scale_shape)
    if self.mode == 'fan_in':
      scale /= max(1., fan_in)
    elif self.mode == 'fan_out':
      scale /= max(1., fan_out)
    else:
      scale /= max(1., (fan_in + fan_out) / 2.)
    if self.distribution == 'truncated_normal':
      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
      stddev = math.sqrt(scale) / .87962566103423978
    else:  # self.distribution == 'untruncated_normal':
      stddev = math.sqrt(scale)
    return tf.random.truncated_normal(shape, mean=stddev, stddev=stddev*0.1,
                                      dtype=dtype, seed=self.seed)


class TrainableHalfCauchy(tf.keras.layers.Layer):
  """Half-Cauchy distribution initializer with trainable parameters."""

  def __init__(self,
               loc_initializer=tf.keras.initializers.TruncatedNormal(
                   stddev=1e-5),
               scale_initializer=tf.keras.initializers.TruncatedNormal(
                   mean=-3., stddev=0.1),
               loc_regularizer=None,
               scale_regularizer=None,
               loc_constraint=None,
               scale_constraint='softplus',
               seed=None,
               **kwargs):
    """Constructs the initializer."""
    super(TrainableHalfCauchy, self).__init__(**kwargs)
    self.loc_initializer = get(loc_initializer)
    self.scale_initializer = get(scale_initializer)
    self.loc_regularizer = regularizers.get(loc_regularizer)
    self.scale_regularizer = regularizers.get(scale_regularizer)
    self.loc_constraint = constraints.get(loc_constraint)
    self.scale_constraint = constraints.get(scale_constraint)
    self.seed = seed

  def build(self, shape, dtype=None):
    if dtype is None:
      dtype = self.dtype

    self.loc = self.add_weight(
        'loc',
        shape=shape,
        initializer=self.loc_initializer,
        regularizer=self.loc_regularizer,
        constraint=None,
        dtype=dtype,
        trainable=True)
    self.scale = self.add_weight(
        'scale',
        shape=shape,
        initializer=self.scale_initializer,
        regularizer=self.scale_regularizer,
        constraint=None,
        dtype=dtype,
        trainable=True)
    self.built = True

  def __call__(self, shape, dtype=None):
    if not self.built:
      self.build(shape, dtype)
    loc = self.loc
    if self.loc_constraint:
      loc = self.loc_constraint(loc)
    scale = self.scale
    if self.scale_constraint:
      scale = self.scale_constraint(scale)
    return generated_random_variables.Independent(
        generated_random_variables.HalfCauchy(loc=loc,
                                              scale=scale).distribution,
        reinterpreted_batch_ndims=len(shape))

  def get_config(self):
    return {
        'loc_initializer':
            serialize(self.loc_initializer),
        'scale_initializer':
            serialize(self.scale_initializer),
        'loc_regularizer':
            regularizers.serialize(self.loc_regularizer),
        'scale_regularizer':
            regularizers.serialize(self.scale_regularizer),
        'loc_constraint':
            constraints.serialize(self.loc_constraint),
        'scale_constraint':
            constraints.serialize(self.scale_constraint),
        'seed': self.seed,
    }


class TrainableNormal(tf.keras.layers.Layer):
  """Random normal op as an initializer with trainable mean and stddev."""

  def __init__(self,
               mean_initializer=tf.keras.initializers.TruncatedNormal(
                   stddev=1e-5),
               stddev_initializer=tf.keras.initializers.TruncatedNormal(
                   mean=-3., stddev=0.1),
               mean_regularizer=None,
               stddev_regularizer=None,
               mean_constraint=None,
               stddev_constraint='softplus',
               seed=None,
               **kwargs):
    """Constructs the initializer."""
    super(TrainableNormal, self).__init__(**kwargs)
    self.mean_initializer = get(mean_initializer)
    self.stddev_initializer = get(stddev_initializer)
    self.mean_regularizer = regularizers.get(mean_regularizer)
    self.stddev_regularizer = regularizers.get(stddev_regularizer)
    self.mean_constraint = constraints.get(mean_constraint)
    self.stddev_constraint = constraints.get(stddev_constraint)
    self.seed = seed

  def build(self, shape, dtype=None):
    if dtype is None:
      dtype = self.dtype

    self.mean = self.add_weight(
        'mean',
        shape=shape,
        initializer=self.mean_initializer,
        regularizer=self.mean_regularizer,
        constraint=None,
        dtype=dtype,
        trainable=True)
    self.stddev = self.add_weight(
        'stddev',
        shape=shape,
        initializer=self.stddev_initializer,
        regularizer=self.stddev_regularizer,
        constraint=None,
        dtype=dtype,
        trainable=True)
    self.built = True

  def __call__(self, shape, dtype=None):
    if not self.built:
      self.build(shape, dtype)
    mean = self.mean
    if self.mean_constraint:
      mean = self.mean_constraint(mean)
    stddev = self.stddev
    if self.stddev_constraint:
      stddev = self.stddev_constraint(stddev)
    return generated_random_variables.Independent(
        generated_random_variables.Normal(loc=mean, scale=stddev).distribution,
        reinterpreted_batch_ndims=len(shape))

  def get_config(self):
    return {
        'mean_initializer':
            serialize(self.mean_initializer),
        'stddev_initializer':
            serialize(self.stddev_initializer),
        'mean_regularizer':
            regularizers.serialize(self.mean_regularizer),
        'stddev_regularizer':
            regularizers.serialize(self.stddev_regularizer),
        'mean_constraint':
            constraints.serialize(self.mean_constraint),
        'stddev_constraint':
            constraints.serialize(self.stddev_constraint),
        'seed': self.seed,
    }


class TrainableHeNormal(TrainableNormal):
  """Trainable normal initialized per He et al. 2015, given a ReLU nonlinearity.

  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in)`,
  where `fan_in` is the number of input units. A ReLU nonlinearity is assumed
  for this initialization scheme.

  References:
    He K, Zhang X, Ren S, Sun J. Delving deep into rectifiers: Surpassing
    human-level performance on imagenet classification. In Proceedings of the
    IEEE international conference on computer vision 2015 (pp. 1026-1034).
    https://arxiv.org/abs/1502.01852
  """

  def __init__(self, seed=None):
    super(TrainableHeNormal, self).__init__(
        mean_initializer=tf.keras.initializers.he_normal(seed),
        seed=seed)

  def get_config(self):
    return {
        'seed': self.seed,
    }


class TrainableGlorotNormal(TrainableNormal):
  """Trainable normal initialized per Glorot and Bengio, 2010.

  The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in +
  fan_out)`, where `fan_in` is the number of input units and `fan_out` is the
  number of output units.

  References:
    Glorot X, Bengio Y. Understanding the difficulty of training deep
    feedforward neural networks. In Proceedings of the thirteenth international
    conference on artificial intelligence and statistics 2010 Mar 31 (pp.
    249-256). http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf
  """

  def __init__(self, seed=None):
    super(TrainableGlorotNormal, self).__init__(
        mean_initializer=tf.keras.initializers.GlorotNormal(seed),
        seed=seed)

  def get_config(self):
    return {
        'seed': self.seed,
    }


class RandomSign(tf.keras.initializers.Initializer):
  """Initializer that generates tensors initialized to +/- 1.

  Attributes:
    probs: probability of +1.
    dtype: tensorflow dtype.
    seed: A Python integer. Used to create random seeds. See
      `tf.set_random_seed`
  """

  def __init__(self, probs=1.0, seed=None, dtype=tf.float32):
    self.probs = probs
    self.seed = seed
    self.dtype = dtype

  def __call__(self, shape, dtype=None, partition_info=None):
    if dtype is None:
      dtype = self.dtype
    bernoulli = tfp.distributions.Bernoulli(probs=self.probs,
                                            dtype=dtype)
    return 2. * bernoulli.sample(shape, self.seed) - 1.

  def get_config(self):
    return {
        'dtype': self.dtype.name,
        'seed': self.seed,
        'probs': self.probs
    }


# Compatibility aliases, following tf.keras

# pylint: disable=invalid-name
scaled_normal_std_dev = ScaledNormalStdDev
trainable_half_cauchy = TrainableHalfCauchy
trainable_normal = TrainableNormal
trainable_he_normal = TrainableHeNormal
trainable_glorot_normal = TrainableGlorotNormal
random_sign = RandomSign
# pylint: enable=invalid-name

# Utility functions, following tf.keras


def serialize(initializer):
  return tf.keras.utils.serialize_keras_object(initializer)


def deserialize(config, custom_objects=None):
  return tf.keras.utils.deserialize_keras_object(
      config,
      module_objects=globals(),
      custom_objects=custom_objects,
      printable_module_name='initializers')


def get(identifier, value=None):
  """Getter for loading from strings; falls back to Keras as needed."""
  if value is None:
    value = identifier
  if identifier is None:
    return None
  elif isinstance(identifier, dict):
    try:
      return deserialize(identifier)
    except ValueError:
      pass
  elif isinstance(identifier, six.string_types):
    config = {'class_name': str(identifier), 'config': {}}
    try:
      return deserialize(config)
    except ValueError:
      pass
  elif callable(identifier):
    return identifier
  return tf.keras.initializers.get(value)

tensorflow.compat.v2.keras.initializers.TruncatedNormal tensorflow.compat.v2.random.truncated_normal tensorflow.compat.v2.keras.utils.serialize_keras_object tensorflow.compat.v2.keras.initializers.he_normal tensorflow.compat.v2.keras.initializers.GlorotNormal tensorflow.compat.v2.keras.initializers.get