# Copyright 2018 The TensorFlow Probability 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.
# ============================================================================
"""Generalized Linear Model Fisher Scoring."""

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

import numpy as np

import tensorflow as tf
from tensorflow_probability.python.glm.util import common_dtype


__all__ = [
    'fit',
    'fit_one_step',
    'convergence_criteria_small_relative_norm_weights_change',
]


def fit(
    model_matrix,
    response,
    model,
    model_coefficients_start=None,
    predicted_linear_response_start=None,
    dispersion=None,
    offset=None,
    convergence_criteria_fn=None,
    learning_rate=None,
    fast_unsafe_numerics=True,
    maximum_iterations=None,
    name=None):
  """Runs multiple Fisher scoring steps.

  Args:
    model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
      represents a sample's features.
    response: (Batch of) vector-shaped `Tensor` where each element represents a
      sample's observed response (to the corresponding row of features). Must
      have same `dtype` as `model_matrix`.
    model: `tfp.glm.ExponentialFamily`-like instance which implicitly
      characterizes a negative log-likelihood loss by specifying the
      distribuion's `mean`, `gradient_mean`, and `variance`.
    model_coefficients_start: Optional (batch of) vector-shaped `Tensor`
      representing the initial model coefficients, one for each column in
      `model_matrix`. Must have same `dtype` as `model_matrix`.
      Default value: `tf.zeros(tf.shape(model_matrix)[-1], model_matrix.dtype)`.
    predicted_linear_response_start: Optional `Tensor` with `shape`, `dtype`
      matching `response`; represents `offset` shifted initial linear
      predictions based on `model_coefficients_start`.
      Default value: `offset` if `model_coefficients is None`, and
      `tf.matmul(model_matrix, model_coefficients_start) + offset` otherwise.
    dispersion: Optional `Tensor` representing `response` dispersion, i.e., as
      in, `p(y|theta) := exp((y theta - A(theta)) / dispersion)`. Must broadcast
      with rows of `model_matrix`.
      Default value: `None` (i.e., "no dispersion").
    offset: Optional `Tensor` with `shape`, `dtype` matching `response`;
      represents constant shift applied to `predicted_linear_response`.
      Default value: `None` (i.e., `tf.zeros_like(response)`).
    convergence_criteria_fn: Python `callable` taking:
      `is_converged_previous`, `iter_`, `model_coefficients_previous`,
      `predicted_linear_response_previous`, `model_coefficients_next`,
      `predicted_linear_response_next`, `response`, `model`, `dispersion` and
      returning a `bool` `Tensor` indicating that Fisher scoring has converged.
      See `convergence_criteria_small_relative_norm_weights_change` as an
      example function.
      Default value: `None` (i.e.,
      `convergence_criteria_small_relative_norm_weights_change`).
    learning_rate: Optional scalar `Tensor` used to dampen iterative progress.
      Typically only needed if optimization diverges, should be no larger than
      `1` and typically very close to `1`.
      Default value: `None` (i.e., `1`).
    fast_unsafe_numerics: Optional Python `bool` indicating if faster, less
      numerically accurate methods can be employed for computing the weighted
      least-squares solution.
      Default value: `True` (i.e., "fast but possibly diminished accuracy").
    maximum_iterations: Optional maximum number of iterations of Fisher scoring
      to run; "and-ed" with result of `convergence_criteria_fn`.
      Default value: `None` (i.e., `infinity`).
    name: Python `str` used as name prefix to ops created by this function.
      Default value: `"fit"`.

  Returns:
    model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
      fitted model coefficients, one for each column in `model_matrix`.
    predicted_linear_response: `response`-shaped `Tensor` representing linear
      predictions based on new `model_coefficients`, i.e.,
      `tf.matmul(model_matrix, model_coefficients) + offset`.
    is_converged: `bool` `Tensor` indicating that the returned
      `model_coefficients` met the `convergence_criteria_fn` criteria within the
      `maximum_iterations` limit.
    iter_: `int32` `Tensor` indicating the number of iterations taken.

  #### Example

  ```python
  from __future__ import print_function
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tf.contrib.distributions

  def make_dataset(n, d, link, scale=1., dtype=np.float32):
    model_coefficients = tfd.Uniform(
        low=np.array(-1, dtype),
        high=np.array(1, dtype)).sample(d, seed=42)
    radius = np.sqrt(2.)
    model_coefficients *= radius / tf.linalg.norm(model_coefficients)
    model_matrix = tfd.Normal(
        loc=np.array(0, dtype),
        scale=np.array(1, dtype)).sample([n, d], seed=43)
    scale = tf.convert_to_tensor(scale, dtype)
    linear_response = tf.tensordot(
        model_matrix, model_coefficients, axes=[[1], [0]])
    if link == 'linear':
      response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
    elif link == 'probit':
      response = tf.cast(
          tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
          dtype)
    elif link == 'logit':
      response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
    else:
      raise ValueError('unrecognized true link: {}'.format(link))
    return model_matrix, response, model_coefficients

  X, Y, w_true = make_dataset(n=int(1e6), d=100, link='probit')

  w, linear_response, is_converged, num_iter = tfp.glm.fit(
      model_matrix=X,
      response=Y,
      model=tfp.glm.BernoulliNormalCDF())
  log_likelihood = tfp.glm.BernoulliNormalCDF().log_prob(Y, linear_response)

  with tf.Session() as sess:
    [w_, linear_response_, is_converged_, num_iter_, Y_, w_true_,
     log_likelihood_] = sess.run([
        w, linear_response, is_converged, num_iter, Y, w_true,
        log_likelihood_])

  print('is_converged: ', is_converged_)
  print('    num_iter: ', num_iter_)
  print('    accuracy: ', np.mean((linear_response_ > 0.) == Y_))
  print('    deviance: ', 2. * np.mean(log_likelihood_))
  print('||w0-w1||_2 / (1+||w0||_2): ', (np.linalg.norm(w_true_ - w_, ord=2) /
                                         (1. + np.linalg.norm(w_true_, ord=2))))

  # ==>
  # is_converged:  True
  #     num_iter:  6
  #     accuracy:  0.804382
  #     deviance:  -0.820746600628
  # ||w0-w1||_2 / (1+||w0||_2):  0.00619245105309
  ```

  """
  graph_deps = [model_matrix, response, model_coefficients_start,
                predicted_linear_response_start, dispersion, offset,
                learning_rate, maximum_iterations]
  with tf.name_scope(name, 'fit', graph_deps):
    [
        model_matrix,
        response,
        model_coefficients_start,
        predicted_linear_response_start,
    ] = prepare_args(
        model_matrix,
        response,
        model_coefficients_start,
        predicted_linear_response_start,
        offset)
    if convergence_criteria_fn is None:
      convergence_criteria_fn = (
          convergence_criteria_small_relative_norm_weights_change())

    def _body(
        is_converged_previous,
        iter_,
        model_coefficients_previous,
        predicted_linear_response_previous):
      """`tf.while_loop` body."""
      model_coefficients_next, predicted_linear_response_next = fit_one_step(
          model_matrix,
          response,
          model,
          predicted_linear_response_previous,
          dispersion,
          offset,
          learning_rate,
          fast_unsafe_numerics)
      is_converged_next = convergence_criteria_fn(
          is_converged_previous=is_converged_previous,
          iter_=iter_,
          model_coefficients_previous=model_coefficients_previous,
          predicted_linear_response_previous=predicted_linear_response_previous,
          model_coefficients_next=model_coefficients_next,
          predicted_linear_response_next=predicted_linear_response_next,
          response=response,
          model=model,
          dispersion=dispersion)
      return [
          is_converged_next,
          iter_ + 1,
          model_coefficients_next,
          predicted_linear_response_next,
      ]

    # while not converged:
    #   fit_one_step
    [
        is_converged,
        iter_,
        model_coefficients,
        predicted_linear_response,
    ] = tf.while_loop(
        cond=lambda is_converged, *args: tf.logical_not(is_converged),
        body=_body,
        loop_vars=[
            tf.zeros([], np.bool),   # is_converged
            tf.zeros([], np.int32),  # iter_
            model_coefficients_start,
            predicted_linear_response_start,
        ],
        maximum_iterations=maximum_iterations)

    # Drop extra dim and set shape hints.
    model_coefficients = model_coefficients[..., 0]
    predicted_linear_response = predicted_linear_response[..., 0]

    return [
        model_coefficients,
        predicted_linear_response,
        is_converged,
        iter_
    ]


def fit_one_step(
    model_matrix,
    response,
    model,
    predicted_linear_response_start,
    dispersion=None,
    offset=None,
    learning_rate=None,
    fast_unsafe_numerics=True,
    name=None):
  """Runs one step of Fisher scoring.

  Args:
    model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
      represents a sample's features.
    response: (Batch of) vector-shaped `Tensor` where each element represents a
      sample's observed response (to the corresponding row of features). Must
      have same `dtype` as `model_matrix`.
    model: `tfp.glm.ExponentialFamily`-like instance used to construct the
      negative log-likelihood loss, gradient, and expected Hessian (i.e., the
      Fisher information matrix).
    predicted_linear_response_start: Optional `Tensor` with `shape`, `dtype`
      matching `response`; represents `offset` shifted initial linear
      predictions based on `model_coefficients_start`.
      Default value: `offset` if `model_coefficients is None`, and
      `tf.matmul(model_matrix, model_coefficients_start) + offset` otherwise.
    dispersion: Optional `Tensor` representing `response` dispersion, i.e., as
      in, `p(y|theta) := exp((y theta - A(theta)) / dispersion)`. Must broadcast
      with rows of `model_matrix`.
      Default value: `None` (i.e., "no dispersion").
    offset: Optional `Tensor` with `shape`, `dtype` matching `response`;
      represents constant shift applied to `predicted_linear_response`.
      Default value: `None` (i.e., `tf.zeros_like(response)`).
    learning_rate: Optional scalar `Tensor` used to dampen iterative progress.
      Typically only needed if optimization diverges, should be no larger than
      `1` and typically very close to `1`.
      Default value: `None` (i.e., `1`).
    fast_unsafe_numerics: Optional Python `bool` indicating if solve should be
      based on Cholesky or QR decomposition.
      Default value: `True` (i.e., "prefer speed via Cholesky decomposition").
    name: Python `str` used as name prefix to ops created by this function.
      Default value: `"fit_one_step"`.

  Returns:
    model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
      fitted model coefficients, one for each column in `model_matrix`.
    predicted_linear_response: `response`-shaped `Tensor` representing linear
      predictions based on new `model_coefficients`, i.e.,
      `tf.matmul(model_matrix, model_coefficients_next) + offset`.
  """
  graph_deps = [model_matrix, response, predicted_linear_response_start,
                dispersion, learning_rate]
  with tf.name_scope(name, 'fit_one_step', graph_deps):
    dtype = common_dtype(graph_deps, np.float32)
    model_matrix = tf.convert_to_tensor(
        model_matrix, dtype=dtype, name='model_matrix')
    response = tf.convert_to_tensor(response, dtype=dtype, name='response')
    predicted_linear_response_start = tf.convert_to_tensor(
        predicted_linear_response_start,
        dtype=dtype,
        name='predicted_linear_response_start')
    # Compute: mean, grad(mean, predicted_linear_response_start), and variance.
    mean, variance, grad_mean = model(predicted_linear_response_start)

    # If either `grad_mean` or `variance is non-finite or zero, then we'll
    # replace it with a value such that the row is zeroed out. Although this
    # procedure may seem circuitous, it is necessary to ensure this algorithm is
    # itself differentiable.
    is_valid = (tf.is_finite(grad_mean) & tf.not_equal(grad_mean, 0.) &
                tf.is_finite(variance) & (variance > 0.))
    def mask_if_invalid(x, mask):
      mask = tf.fill(tf.shape(x), value=np.array(mask, x.dtype.as_numpy_dtype))
      return tf.where(is_valid, x, mask)

    # Run one step of iteratively reweighted least-squares.
    # Compute "`z`", the adjusted predicted linear response.
    # z = predicted_linear_response_start
    #     + learning_rate * (response - mean) / grad_mean
    z = (response - mean) / mask_if_invalid(grad_mean, 1.)
    # TODO(jvdillon): Rather than use learning rate, we should consider using
    # backtracking line search.
    if learning_rate is not None:
      z *= learning_rate
    z += predicted_linear_response_start

    # Compute "`w`", the per-sample weight.
    if dispersion is not None:
      # For convenience, we'll now scale the variance by the dispersion factor.
      variance *= dispersion
    w = (mask_if_invalid(grad_mean, 0.) *
         tf.rsqrt(mask_if_invalid(variance, np.inf)))

    a = model_matrix * w
    b = z * w
    # Solve `min{ || A @ model_coefficients - b ||_2**2 : model_coefficients }`
    # where `@` denotes `matmul`.
    model_coefficients_next = tf.matrix_solve_ls(
        a, b, fast=fast_unsafe_numerics, name='model_coefficients_next')

    # TODO(b/79122261): The approach used in `matrix_solve_ls` could be made
    # faster by avoiding explicitly forming Q and instead keeping the
    # factorization in 'implicit' form with stacked (rescaled) Householder
    # vectors underneath the 'R' and then applying the (accumulated)
    # reflectors in the appropriate order to apply Q'. However, we don't
    # presently do this because we lack core TF functionality. For reference,
    # the vanilla QR approach is:
    #   q, r = tf.linalg.qr(a)
    #   c = tf.matmul(q, b, adjoint_a=True)
    #   model_coefficients_next = tf.matrix_triangular_solve(
    #       r, c, lower=False, name='model_coefficients_next')
    predicted_linear_response_next = tf.matmul(
        model_matrix,
        model_coefficients_next,
        name='predicted_linear_response_next')

    if offset is not None:
      predicted_linear_response_next += offset

    return model_coefficients_next, predicted_linear_response_next


def convergence_criteria_small_relative_norm_weights_change(
    tolerance=1e-5,
    norm_order=2):
  """Returns Python `callable` which indicates fitting procedure has converged.

  Writing old, new `model_coefficients` as `w0`, `w1`, this function
  defines convergence as,

  ```python
  relative_euclidean_norm = (tf.norm(w0 - w1, ord=2, axis=-1) /
                             (1. + tf.norm(w0, ord=2, axis=-1)))
  reduce_all(relative_euclidean_norm < tolerance)
  ```

  where `tf.norm(x, ord=2)` denotes the [Euclidean norm](
  https://en.wikipedia.org/wiki/Norm_(mathematics)#Euclidean_norm) of `x`.

  Args:
    tolerance: `float`-like `Tensor` indicating convergence, i.e., when
      max relative Euclidean norm weights difference < tolerance`.
      Default value: `1e-5`.
    norm_order: Order of the norm. Default value: `2` (i.e., "Euclidean norm".)

  Returns:
    convergence_criteria_fn: Python `callable` which returns `bool` `Tensor`
      indicated fitting procedure has converged. (See inner function
      specification for argument signature.)
      Default value: `1e-5`.
  """
  def convergence_criteria_fn(
      is_converged_previous,  # pylint: disable=unused-argument
      iter_,
      model_coefficients_previous,
      predicted_linear_response_previous,  # pylint: disable=unused-argument
      model_coefficients_next,
      predicted_linear_response_next,  # pylint: disable=unused-argument
      response,  # pylint: disable=unused-argument
      model,  # pylint: disable=unused-argument
      dispersion):  # pylint: disable=unused-argument
    """Returns `bool` `Tensor` indicating if fitting procedure has converged.

    Args:
      is_converged_previous: "old" convergence results.
      iter_: Iteration number.
      model_coefficients_previous: "old" `model_coefficients`.
      predicted_linear_response_previous: "old" `predicted_linear_response`.
      model_coefficients_next: "new" `model_coefficients`.
      predicted_linear_response_next: "new: `predicted_linear_response`.
      response: (Batch of) vector-shaped `Tensor` where each element represents
        a sample's observed response (to the corresponding row of features).
        Must have same `dtype` as `model_matrix`.
      model: `tfp.glm.ExponentialFamily`-like instance used to construct the
        negative log-likelihood loss, gradient, and expected Hessian (i.e., the
        Fisher information matrix).
      dispersion: `Tensor` representing `response` dispersion, i.e., as in:
        `p(y|theta) := exp((y theta - A(theta)) / dispersion)`. Must broadcast
        with rows of `model_matrix`.
        Default value: `None` (i.e., "no dispersion").

    Returns:
      is_converged: `bool` `Tensor`.
    """
    relative_euclidean_norm = (
        tf.norm(model_coefficients_previous -
                model_coefficients_next, ord=norm_order, axis=-1)
        / (1. + tf.norm(model_coefficients_previous, ord=norm_order, axis=-1)))
    return (iter_ > 0) & tf.reduce_all(relative_euclidean_norm < tolerance)
  return convergence_criteria_fn


def prepare_args(model_matrix,
                 response,
                 model_coefficients,
                 predicted_linear_response,
                 offset,
                 name=None):
  """Helper to `fit` which sanitizes input args.

  Args:
    model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
      represents a sample's features.
    response: (Batch of) vector-shaped `Tensor` where each element represents a
      sample's observed response (to the corresponding row of features). Must
      have same `dtype` as `model_matrix`.
    model_coefficients: Optional (batch of) vector-shaped `Tensor` representing
      the model coefficients, one for each column in `model_matrix`. Must have
      same `dtype` as `model_matrix`.
      Default value: `tf.zeros(tf.shape(model_matrix)[-1], model_matrix.dtype)`.
    predicted_linear_response: Optional `Tensor` with `shape`, `dtype` matching
      `response`; represents `offset` shifted initial linear predictions based
      on current `model_coefficients`.
      Default value: `offset` if `model_coefficients is None`, and
      `tf.matmul(model_matrix, model_coefficients_start) + offset` otherwise.
    offset: Optional `Tensor` with `shape`, `dtype` matching `response`;
      represents constant shift applied to `predicted_linear_response`.
      Default value: `None` (i.e., `tf.zeros_like(response)`).
    name: Python `str` used as name prefix to ops created by this function.
      Default value: `"prepar_args"`.

  Returns:
    model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
      fitted model coefficients, one for each column in `model_matrix`.
    predicted_linear_response: `response`-shaped `Tensor` representing linear
      predictions based on `model_coefficients`, i.e.,
      `tf.matmul(model_matrix, model_coefficients) + offset`.
  """
  graph_deps = [model_matrix, response, model_coefficients,
                predicted_linear_response, offset]
  with tf.name_scope(name, 'prepare_args', graph_deps):
    dtype = common_dtype(graph_deps, np.float32)

    model_matrix = tf.convert_to_tensor(
        model_matrix, dtype=dtype, name='model_matrix')

    response = tf.convert_to_tensor(response, dtype=dtype, name='response')
    response = response[..., tf.newaxis]  # [B, 1]

    use_default_model_coefficients = model_coefficients is None
    if use_default_model_coefficients:
      # User did not supply model coefficients; assume they're all zero.
      model_coefficients = tf.zeros(
          [num_cols(model_matrix)], dtype=dtype, name='model_coefficients')
    else:
      # User did supply model coefficients; convert to Tensor in case it's
      # numpy or literal.
      model_coefficients = tf.convert_to_tensor(
          model_coefficients, dtype=dtype, name='model_coefficients')

    model_coefficients = model_coefficients[..., tf.newaxis]  # [B, d, 1]

    if use_default_model_coefficients:
      # Since we're using zeros for model_coefficients, we know the predicted
      # linear response will also be all zeros.
      if (model_matrix.shape.ndims is not None and
          model_matrix.shape[:-1].is_fully_defined()):
        # Obtain shape statically.
        shape = np.int32(model_matrix.shape[:-1].concatenate([1]).as_list())
      else:
        # Obtain shape dynamically.
        shape = tf.concat([tf.shape(model_matrix)[:-1], [1]], axis=0)
      predicted_linear_response = tf.zeros(
          shape, dtype, name='predicted_linear_response')
      if offset is not None:
        predicted_linear_response += offset
    elif predicted_linear_response is None:
      # We were given model_coefficiencts but not the predicted linear response,
      # hence compute it (as a matmul).
      predicted_linear_response = tf.matmul(
          model_matrix, model_coefficients, name='predicted_linear_response')
      if offset is not None:
        predicted_linear_response += offset
    else:
      # We were given the predicted linear response. However, we assume the user
      # supplies a vector not a `n x 1` matrix. For their convenience we expand
      # the `[n]`-shaped Tensor to `[n, 1]`.
      predicted_linear_response = tf.expand_dims(
          predicted_linear_response, axis=-1, name='predicted_linear_response')

    return model_matrix, response, model_coefficients, predicted_linear_response


def num_cols(x):
  """Returns number of cols in a given `Tensor`."""
  if x.shape.ndims is not None and x.shape[-1].value is not None:
    return x.shape[-1].value
  return tf.shape(x)[-1]
