# Copyright 2020 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.
# ============================================================================
"""Tests for multivariate Power Spherical distribution."""

# Dependency imports
import numpy as np
from scipy import special as sp_special

import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp

from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import tensorshape_util
from tensorflow_probability.python.internal import test_util


class _PowerSphericalTest(object):

  def testReproducibleGraph(self):
    pspherical = tfp.distributions.PowerSpherical(
        mean_direction=tf.math.l2_normalize(np.array(
            [1., 2.], dtype=self.dtype)),
        concentration=self.dtype(1.2))
    seed = test_util.test_seed()
    s1 = self.evaluate(pspherical.sample(50, seed=seed))
    if tf.executing_eagerly():
      tf.random.set_seed(seed)
    s2 = self.evaluate(pspherical.sample(50, seed=seed))
    self.assertAllEqual(s1, s2)

  def VerifySampleMean(self, mean_dirs, concentration, batch_shape):
    pspherical = tfp.distributions.PowerSpherical(
        mean_direction=mean_dirs,
        concentration=concentration,
        validate_args=True,
        allow_nan_stats=False)
    self.assertEqual([mean_dirs.shape[-1]],
                     tensorshape_util.as_list(pspherical.event_shape))
    self.assertEqual(
        batch_shape,
        tensorshape_util.as_list(pspherical.batch_shape))
    nsamples = int(5e4)
    samples = pspherical.sample(nsamples, seed=test_util.test_seed())
    self.assertEqual([nsamples] + batch_shape + [mean_dirs.shape[-1]],
                     tensorshape_util.as_list(samples.shape))
    sample_mean = self.evaluate(samples).mean(axis=0)
    sample_dir = (
        sample_mean / np.linalg.norm(sample_mean, axis=-1, keepdims=True))
    # Assert that positive-concentration distributions have samples with
    # the expected mean direction.
    inner_product = self.evaluate(
        tf.reduce_sum(sample_dir * pspherical.mean_direction, axis=-1))
    # Inner products should be roughly ascending by concentration.
    self.assertAllClose(np.round(np.sort(inner_product, axis=0), decimals=3),
                        np.round(inner_product, decimals=3),
                        atol=.007)
    means = self.evaluate(pspherical.mean())
    # Mean vector for 0-concentration is precisely (0, 0).
    self.assertAllEqual(np.zeros_like(means[0]), means[0])
    mean_lengths = np.linalg.norm(means, axis=-1)
    # Length of the mean vector is strictly ascending with concentration.
    self.assertAllEqual(mean_lengths, np.sort(mean_lengths, axis=0))
    self.assertAllClose(np.linalg.norm(sample_mean, axis=-1), mean_lengths,
                        atol=0.03)

  def testSampleMeanDir2d(self):
    mean_dirs = tf.math.l2_normalize(
        np.array([[1., 1], [-2, 1], [0, -1]], dtype=self.dtype), axis=-1)
    concentration = np.array(
        [[0], [0.1], [2], [40], [1000]], dtype=self.dtype)
    self.VerifySampleMean(mean_dirs, concentration, [5, 3])

  def testSampleMeanDir3d(self):
    mean_dirs = tf.math.l2_normalize(
        np.array([[1., 2, 3], [-2, -3, -1]], dtype=self.dtype), axis=-1)
    concentration = np.array(
        [[0], [0.1], [2], [40], [1000]], dtype=self.dtype)
    self.VerifySampleMean(mean_dirs, concentration, [5, 2])

  def testSampleMeanDir5d(self):
    mean_dirs = tf.math.l2_normalize(
        np.array([[1., 2, 3, -1., 5.]], dtype=self.dtype), axis=-1)
    concentration = np.array(
        [[0], [0.1], [2], [40], [1000]], dtype=self.dtype)
    self.VerifySampleMean(mean_dirs, concentration, [5, 1])

  def VerifyPdfWithNumpy(self, pspherical, atol=1e-4):
    """Verifies log_prob evaluations with numpy/scipy.

    Both uniform random points and sampled points are evaluated.

    Args:
      pspherical: A `tfp.distributions.PowerSpherical` instance.
      atol: Absolute difference tolerable.
    """
    dim = tf.compat.dimension_value(pspherical.event_shape[-1])
    nsamples = 10
    # Sample some random points uniformly over the hypersphere using numpy.
    sample_shape = [nsamples] + tensorshape_util.as_list(
        pspherical.batch_shape) + [dim]
    uniforms = np.random.randn(*sample_shape)
    uniforms /= np.linalg.norm(uniforms, axis=-1, keepdims=True)
    uniforms = uniforms.astype(dtype_util.as_numpy_dtype(pspherical.dtype))
    # Concatenate in some sampled points from the distribution under test.
    pspherical_samples = pspherical.sample(
        sample_shape=[nsamples], seed=test_util.test_seed())
    samples = tf.concat([uniforms, pspherical_samples], axis=0)
    samples = tf.debugging.check_numerics(samples, 'samples')
    samples = self.evaluate(samples)
    log_prob = pspherical.log_prob(samples)
    log_prob = self.evaluate(log_prob)

    # Check that the log_prob is not nan or +inf. It can be -inf since
    # if we sample a direction diametrically opposite to the mean direction,
    # we'll get an inner product of -1.
    self.assertFalse(np.any(np.isnan(log_prob)))
    self.assertFalse(np.any(np.isposinf(log_prob)))
    conc = self.evaluate(pspherical.concentration)
    mean_dir = self.evaluate(pspherical.mean_direction)
    alpha = (dim - 1.) / 2. + conc
    beta = (dim - 1.) / 2.

    expected = (
        sp_special.xlog1py(conc, np.sum(samples * mean_dir, axis=-1)) -
        (alpha + beta) * np.log(2.) - beta * np.log(np.pi) -
        sp_special.gammaln(alpha) + sp_special.gammaln(alpha + beta))
    self.assertAllClose(expected, log_prob, atol=atol)

  def VerifySampleAndPdfConsistency(self, pspherical, rtol=0.075):
    """Verifies samples are consistent with the PDF using importance sampling.

    In particular, we verify an estimate the surface area of the n-dimensional
    hypersphere, and the surface areas of the spherical caps demarcated by
    a handful of survival rates.

    Args:
      pspherical: A `PowerSpherical` distribution instance.
      rtol: Relative difference tolerable.
    """
    dim = tf.compat.dimension_value(pspherical.event_shape[-1])
    nsamples = int(1e5)
    samples = pspherical.sample(
        sample_shape=[nsamples], seed=test_util.test_seed())
    samples = tf.debugging.check_numerics(samples, 'samples')
    log_prob = pspherical.log_prob(samples)
    log_prob, samples = self.evaluate([log_prob, samples])
    # Check that the log_prob is not nan or +inf. It can be -inf since
    # if we sample a direction diametrically opposite to the mean direction,
    # we'll get an inner product of -1.
    self.assertFalse(np.any(np.isnan(log_prob)))
    self.assertFalse(np.any(np.isposinf(log_prob)))
    log_importance = -log_prob
    sphere_surface_area_estimate, importance = self.evaluate([
        tf.reduce_mean(tf.math.exp(log_importance), axis=0),
        tf.exp(log_importance)])
    true_sphere_surface_area = 2 * (np.pi)**(dim / 2) * self.evaluate(
        tf.exp(-tf.math.lgamma(dim / 2)))
    # Broadcast to correct size
    true_sphere_surface_area += np.zeros_like(sphere_surface_area_estimate)
    # Highly concentrated distributions do not get enough coverage to provide
    # a reasonable full-sphere surface area estimate. These are covered below
    # by CDF-based hypersphere cap surface area estimates.
    # Because the PowerSpherical distribution has zero mass at
    # -`mean_direction` (and points close to -`mean_direction` due to floating
    # point), we only compute this at concentration = 0, which has guaranteed
    # mass everywhere.
    self.assertAllClose(
        true_sphere_surface_area[0],
        sphere_surface_area_estimate[0], rtol=rtol)

    # Assert surface area of hyperspherical cap For some CDFs in [.05,.45],
    # (h must be greater than 0 for the hypersphere cap surface area
    # calculation to hold).
    for survival_rate in 0.95, .9, .75, .6:
      cdf = (1 - survival_rate)
      mean_dir = self.evaluate(pspherical.mean_direction)
      dotprods = np.sum(samples * mean_dir, -1)
      # Empirical estimate of the effective dot-product of the threshold that
      # selects for a given CDF level, that is the cosine of the largest
      # passable angle, or the minimum cosine for a within-CDF sample.
      dotprod_thresh = np.percentile(
          dotprods, 100 * survival_rate, axis=0, keepdims=True)
      # We mask this sum because it is possible for the log_prob to be -inf when
      # the mean_direction is -mean_dir.
      importance_masked = np.ma.array(
          importance, mask=dotprods <= dotprod_thresh)
      sphere_cap_surface_area_ests = (
          cdf * (importance_masked).sum(0) /
          (dotprods > dotprod_thresh).sum(0))
      h = (1 - dotprod_thresh)
      self.assertGreaterEqual(h.min(), 0)  # h must be >= 0 for the eqn below
      true_sphere_cap_surface_area = (
          0.5 * true_sphere_surface_area *
          self.evaluate(tf.math.betainc((dim - 1) / 2, 0.5, 2 * h - h**2)))
      if dim == 3:  # For 3-d we have a simpler form we can double-check.
        self.assertAllClose(2 * np.pi * h, true_sphere_cap_surface_area)

      self.assertAllClose(
          true_sphere_cap_surface_area,
          sphere_cap_surface_area_ests +
          np.zeros_like(true_sphere_cap_surface_area),
          rtol=rtol)

  def testSampleAndPdfConsistency2d(self):
    mean_dir = tf.math.l2_normalize([[1., 2], [-2, -3]], axis=-1)
    concentration = [[0], [1e-5], [0.1], [1], [4]]
    pspherical = tfp.distributions.PowerSpherical(
        mean_direction=mean_dir, concentration=concentration,
        validate_args=True, allow_nan_stats=False)
    self.VerifySampleAndPdfConsistency(pspherical)
    self.VerifyPdfWithNumpy(pspherical)

  def testSampleAndPdfConsistency3d(self):
    mean_dir = tf.math.l2_normalize([[1., 2, 3], [-2, -3, -1]], axis=-1)
    concentration = [[0], [1e-5], [0.1], [1], [4]]
    pspherical = tfp.distributions.PowerSpherical(
        mean_direction=mean_dir, concentration=concentration,
        validate_args=True, allow_nan_stats=False)
    self.VerifySampleAndPdfConsistency(pspherical)
    self.VerifyPdfWithNumpy(pspherical, atol=.002)

  def testSampleAndPdfConsistency4d(self):
    mean_dir = tf.math.l2_normalize([[1., 2, 3, 4], [-2, -3, -1, 0]], axis=-1)
    concentration = [[0], [1e-4], [0.1], [1], [4]]
    pspherical = tfp.distributions.PowerSpherical(
        mean_direction=mean_dir, concentration=concentration,
        validate_args=True, allow_nan_stats=False)
    self.VerifySampleAndPdfConsistency(pspherical)
    self.VerifyPdfWithNumpy(pspherical)

  def testSampleAndPdfConsistency5d(self):
    mean_dir = tf.math.l2_normalize(
        [[1., 2, 3, 4, 5], [-2, -3, -1, 0, 1]], axis=-1)
    concentration = [[0], [5e-2], [0.1], [1], [4]]
    pspherical = tfp.distributions.PowerSpherical(
        mean_direction=mean_dir, concentration=concentration,
        validate_args=True, allow_nan_stats=False)
    self.VerifySampleAndPdfConsistency(pspherical)
    self.VerifyPdfWithNumpy(pspherical, atol=2e-4)

  def testSampleAndPdfForMeanDirNorthPole(self):
    mean_dir = np.array([1., 0., 0., 0., 0.], dtype=self.dtype)