# 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.
# ============================================================================
"""Tests for multivariate von Mises-Fisher distribution."""

# Dependency imports
from absl.testing import parameterized
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


@test_util.test_all_tf_execution_regimes
class VonMisesFisherTest(test_util.VectorDistributionTestHelpers,
                         test_util.TestCase):

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

  def testSampleMeanDir2d(self):
    mean_dirs = tf.math.l2_normalize([[1., 1], [-2, 1], [0, -1]], axis=-1)
    concentration = [[0], [0.1], [2], [40], [1000]]
    vmf = tfp.distributions.VonMisesFisher(
        mean_direction=mean_dirs,
        concentration=concentration,
        validate_args=True,
        allow_nan_stats=False)
    self.assertEqual([5, 3], tensorshape_util.as_list(vmf.batch_shape))
    self.assertEqual([2], tensorshape_util.as_list(vmf.event_shape))
    nsamples = int(2e4)
    samples = vmf.sample(nsamples, seed=test_util.test_seed())
    self.assertEqual([nsamples, 5, 3, 2],
                     tensorshape_util.as_list(samples.shape))
    sample_mean = self.evaluate(samples).mean(axis=0)
    # Assert that positive-concentration distributions have samples with
    # the expected mean direction.
    sample_dir = (
        sample_mean / np.linalg.norm(sample_mean, axis=-1, keepdims=True))
    inner_product = self.evaluate(
        tf.reduce_sum(sample_dir * vmf.mean_direction, axis=-1))
    # All except the 0-concentration distribution should have >0 inner product
    # with the mean direction of the distribution.
    self.assertAllGreater(inner_product[1:], 0.1)
    # Pick out >1 concentration distributions to assert ~1 inner product with
    # mean direction.
    self.assertAllClose(np.ones_like(inner_product)[2:], inner_product[2:],
                        atol=1e-3)
    # 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=.005)
    means = self.evaluate(vmf.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 testSampleMeanDir3d(self):
    mean_dir = tf.math.l2_normalize([[1., 2, 3], [-2, -3, -1]], axis=-1)
    concentration = [[0], [0.1], [2], [40], [1000]]
    vmf = tfp.distributions.VonMisesFisher(
        mean_direction=mean_dir,
        concentration=concentration,
        validate_args=True,
        allow_nan_stats=False)
    self.assertEqual([5, 2], tensorshape_util.as_list(vmf.batch_shape))
    self.assertEqual([3], tensorshape_util.as_list(vmf.event_shape))
    nsamples = int(2e4)
    samples = vmf.sample(
        sample_shape=[nsamples], seed=test_util.test_seed())
    self.assertEqual([nsamples, 5, 2, 3],
                     tensorshape_util.as_list(samples.shape))
    sample_mean = self.evaluate(samples).mean(axis=0)
    # Assert that positive-concentration distributions have samples with
    # the expected mean direction.
    sample_dir = (
        sample_mean / np.linalg.norm(sample_mean, axis=-1, keepdims=True))
    inner_product = self.evaluate(
        tf.reduce_sum(sample_dir * vmf.mean_direction, axis=-1))
    # All except the 0-concentration distribution should have >0 inner product
    # with the mean direction of the distribution.
    self.assertAllGreater(inner_product[1:], 0.1)
    # Pick out >1 concentration distributions to assert ~1 inner product with
    # mean direction.
    self.assertAllClose(np.ones_like(inner_product)[2:], inner_product[2:],
                        atol=1e-3)
    # Inner products should be roughly ascending by concentration.
    self.assertAllEqual(np.round(np.sort(inner_product, axis=0), decimals=3),
                        np.round(inner_product, decimals=3))
    means = self.evaluate(vmf.mean())
    # Mean vector for 0-concentration is precisely (0, 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 _verifyPdfWithNumpy(self, vmf, atol=1e-4):
    """Verifies log_prob evaluations with numpy/scipy.

    Both uniform random points and sampled points are evaluated.

    Args:
      vmf: A `tfp.distributions.VonMisesFisher` instance.
      atol: Absolute difference tolerable.
    """
    dim = tf.compat.dimension_value(vmf.event_shape[-1])
    nsamples = 10
    # Sample some random points uniformly over the hypersphere using numpy.
    sample_shape = [nsamples] + tensorshape_util.as_list(
        vmf.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(vmf.dtype))
    # Concatenate in some sampled points from the distribution under test.
    vmf_samples = vmf.sample(
        sample_shape=[nsamples], seed=test_util.test_seed())
    samples = tf.concat([uniforms, vmf_samples], axis=0)
    log_prob = vmf.log_prob(samples)
    samples = tf.debugging.check_numerics(samples, 'samples')
    log_prob = tf.debugging.check_numerics(log_prob, 'log_prob')
    samples, log_prob, conc, mean_dir = self.evaluate([
        samples, log_prob, vmf.concentration, vmf.mean_direction])
    log_true_sphere_surface_area = (
        np.log(2) + (dim / 2) * np.log(np.pi) - sp_special.gammaln(dim / 2))
    expected = (
        conc * np.sum(samples * mean_dir, axis=-1) +
        np.where(conc > 0,
                 (dim / 2 - 1) * np.log(conc) -
                 (dim / 2) * np.log(2 * np.pi) -
                 np.log(sp_special.ive(dim / 2 - 1, conc)) -
                 np.abs(conc),
                 -log_true_sphere_surface_area))
    self.assertAllClose(expected, log_prob, atol=atol)

  def _verifySampleAndPdfConsistency(self, vmf, 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:
      vmf: A `VonMisesFisher` distribution instance.
      rtol: Relative difference tolerable.
    """
    dim = tf.compat.dimension_value(vmf.event_shape[-1])
    nsamples = int(6e4)
    samples = vmf.sample(
        sample_shape=[nsamples], seed=test_util.test_seed())
    samples = tf.debugging.check_numerics(samples, 'samples')
    log_prob = vmf.log_prob(samples)
    log_prob = tf.debugging.check_numerics(log_prob, 'log_prob')
    log_importance = -log_prob
    sphere_surface_area_estimate, samples, importance, conc = self.evaluate([
        tf.exp(
            tf.reduce_logsumexp(log_importance, axis=0) -
            tf.math.log(tf.cast(nsamples, dtype=tf.float32))), samples,
        tf.exp(log_importance), vmf.concentration
    ])
    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.
    self.assertAllClose(
        true_sphere_surface_area[np.where(conc < 3)],
        sphere_surface_area_estimate[np.where(conc < 3)],
        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(vmf.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)
      dotprod_above_thresh = np.float32(dotprods > dotprod_thresh)
      sphere_cap_surface_area_ests = (
          cdf * (importance * dotprod_above_thresh).sum(0) /
          dotprod_above_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 _verifyCovariance(self, vmf):
    dim = tf.compat.dimension_value(vmf.event_shape[-1])
    nsamples = 10000