"""Tests for the ZeroInflated distribution."""
from __future__ import absolute_import, division, print_function

# Dependency imports
import numpy as np

import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow_probability.python.internal import test_util as tfp_test_util
from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
tfd = tfp.distributions

from odin.bay import ZeroInflated

def make_synthetic_data(shape, rate,
      count_distribution=lambda size: np.random.negative_binomial(n=5, p=0.3, size=size)):
  """ Generating samples according to Zero-inflation process
  shape : output shape
  rate : zero-inflation rate
  count_distribution : distribution

  Return
  ------
  data_corrupted : corrupted data by zero inflation
  data : original data
  mask : indicator which entries were dropped out
  """
  data = count_distribution(size=shape)
  mask = np.random.binomial(n=1, p=np.clip(1. - rate, 0., 1.),
                            size=shape)
  data_corrupted = (data * mask)
  return data_corrupted, data, 1 - mask

@test_util.run_all_in_graph_and_eager_modes
class ZeroInflatedTest(tf.test.TestCase):

  def testP(self):
    p = [0.2, 0.4]
    dist = tfd.Bernoulli(probs=p)
    self.assertAllClose(p, self.evaluate(dist.probs))

  def testInvalidP(self):
    invalid_ps = [1.01, 2.]
    for p in invalid_ps:
      with self.assertRaisesOpError("probs has components greater than 1"):
        dist = tfd.Bernoulli(probs=p, validate_args=True)
        self.evaluate(dist.probs)

    invalid_ps = [-0.01, -3.]
    for p in invalid_ps:
      with self.assertRaisesOpError("Condition x >= 0"):
        dist = tfd.Bernoulli(probs=p, validate_args=True)
        self.evaluate(dist.probs)

    valid_ps = [0.0, 0.5, 1.0]
    for p in valid_ps:
      dist = tfd.Bernoulli(probs=p)
      self.assertEqual(p, self.evaluate(dist.probs))  # Should not fail

  def testShapes(self):
    for batch_shape in ([], [1], [2, 3, 4]):
      dist = make_bernoulli(batch_shape)
      self.assertAllEqual(batch_shape, dist.batch_shape.as_list())
      self.assertAllEqual(batch_shape, self.evaluate(dist.batch_shape_tensor()))
      self.assertAllEqual([], dist.event_shape.as_list())
      self.assertAllEqual([], self.evaluate(dist.event_shape_tensor()))

  def testDtype(self):
    dist = make_bernoulli([])
    self.assertEqual(dist.dtype, tf.int32)
    self.assertEqual(dist.dtype, dist.sample(5).dtype)
    self.assertEqual(dist.dtype, dist.mode().dtype)
    self.assertEqual(dist.probs.dtype, dist.mean().dtype)
    self.assertEqual(dist.probs.dtype, dist.variance().dtype)
    self.assertEqual(dist.probs.dtype, dist.stddev().dtype)
    self.assertEqual(dist.probs.dtype, dist.entropy().dtype)
    self.assertEqual(dist.probs.dtype, dist.prob(0).dtype)
    self.assertEqual(dist.probs.dtype, dist.prob(0.5).dtype)
    self.assertEqual(dist.probs.dtype, dist.log_prob(0).dtype)
    self.assertEqual(dist.probs.dtype, dist.log_prob(0.5).dtype)

    dist64 = make_bernoulli([], tf.int64)
    self.assertEqual(dist64.dtype, tf.int64)
    self.assertEqual(dist64.dtype, dist64.sample(5).dtype)
    self.assertEqual(dist64.dtype, dist64.mode().dtype)

  def _testPmf(self, **kwargs):
    dist = tfd.Bernoulli(**kwargs)
    # pylint: disable=bad-continuation
    xs = [
        0,
        [1],
        [1, 0],
        [[1, 0]],
        [[1, 0], [1, 1]],
    ]
    expected_pmfs = [
        [[0.8, 0.6], [0.7, 0.4]],
        [[0.2, 0.4], [0.3, 0.6]],
        [[0.2, 0.6], [0.3, 0.4]],
        [[0.2, 0.6], [0.3, 0.4]],
        [[0.2, 0.6], [0.3, 0.6]],
    ]
    # pylint: enable=bad-continuation

    for x, expected_pmf in zip(xs, expected_pmfs):
      self.assertAllClose(self.evaluate(dist.prob(x)), expected_pmf)
      self.assertAllClose(self.evaluate(dist.log_prob(x)), np.log(expected_pmf))

  def testPmfCorrectBroadcastDynamicShape(self):
    p = tf.compat.v1.placeholder_with_default([0.2, 0.3, 0.4], shape=None)
    dist = tfd.Bernoulli(probs=p)
    event1 = [1, 0, 1]
    event2 = [[1, 0, 1]]
    self.assertAllClose(
        [0.2, 0.7, 0.4], self.evaluate(dist.prob(event1)))
    self.assertAllClose(
        [[0.2, 0.7, 0.4]], self.evaluate(dist.prob(event2)))

  def testPmfInvalid(self):
    p = [0.1, 0.2, 0.7]
    dist = tfd.Bernoulli(probs=p, validate_args=True)
    with self.assertRaisesOpError("must be non-negative."):
      self.evaluate(dist.prob([1, 1, -1]))
    with self.assertRaisesOpError("Elements cannot exceed 1."):
      self.evaluate(dist.prob([2, 0, 1]))

  def testPmfWithP(self):
    p = [[0.2, 0.4], [0.3, 0.6]]
    self._testPmf(probs=p)
    if not special:
      return
    self._testPmf(logits=special.logit(p))

  def testPmfWithFloatArgReturnsXEntropy(self):
    p = [[0.2], [0.4], [0.3], [0.6]]
    samps = [0, 0.1, 0.8]
    self.assertAllClose(
        np.float32(samps) * np.log(np.float32(p)) +
        (1 - np.float32(samps)) * np.log(1 - np.float32(p)),
        self.evaluate(
            tfd.Bernoulli(probs=p, validate_args=False).log_prob(samps)))

  def testBroadcasting(self):
    probs = lambda p: tf.compat.v1.placeholder_with_default(p, shape=None)
    dist = lambda p: tfd.Bernoulli(probs=probs(p))
    self.assertAllClose(np.log(0.5), self.evaluate(dist(0.5).log_prob(1)))
    self.assertAllClose(
        np.log([0.5, 0.5, 0.5]), self.evaluate(dist(0.5).log_prob([1, 1, 1])))
    self.assertAllClose(np.log([0.5, 0.5, 0.5]),
                        self.evaluate(dist([0.5, 0.5, 0.5]).log_prob(1)))

  def testPmfShapes(self):
    probs = lambda p: tf.compat.v1.placeholder_with_default(p, shape=None)
    dist = lambda p: tfd.Bernoulli(probs=probs(p))
    self.assertEqual(
        2, len(self.evaluate(dist([[0.5], [0.5]]).log_prob(1)).shape))

    dist = tfd.Bernoulli(probs=0.5)
    self.assertEqual(2, len(self.evaluate(dist.log_prob([[1], [1]])).shape))

    dist = tfd.Bernoulli(probs=0.5)
    self.assertEqual((), dist.log_prob(1).shape)
    self.assertEqual((1), dist.log_prob([1]).shape)
    self.assertEqual((2, 1), dist.log_prob([[1], [1]]).shape)

    dist = tfd.Bernoulli(probs=[[0.5], [0.5]])
    self.assertEqual((2, 1), dist.log_prob(1).shape)

  def testBoundaryConditions(self):
    dist = tfd.Bernoulli(probs=1.0)
    self.assertAllClose(np.nan, self.evaluate(dist.log_prob(0)))
    self.assertAllClose([np.nan], [self.evaluate(dist.log_prob(1))])

  def testEntropyNoBatch(self):
    p = 0.2
    dist = tfd.Bernoulli(probs=p)
    self.assertAllClose(self.evaluate(dist.entropy()), entropy(p))

  def testEntropyWithBatch(self):
    p = [[0.1, 0.7], [0.2, 0.6]]
    dist = tfd.Bernoulli(probs=p, validate_args=False)
    self.assertAllClose(
        self.evaluate(dist.entropy()),
        [[entropy(0.1), entropy(0.7)], [entropy(0.2),
                                        entropy(0.6)]])

  def testSampleN(self):
    p = [0.2, 0.6]
    dist = tfd.Bernoulli(probs=p)
    n = 100000
    samples = dist.sample(n)
    samples.set_shape([n, 2])
    self.assertEqual(samples.dtype, tf.int32)
    sample_values = self.evaluate(samples)
    self.assertTrue(np.all(sample_values >= 0))
    self.assertTrue(np.all(sample_values <= 1))
    # Note that the standard error for the sample mean is ~ sqrt(p * (1 - p) /
    # n). This means that the tolerance is very sensitive to the value of p
    # as well as n.
    self.assertAllClose(p, np.mean(sample_values, axis=0), atol=1e-2)
    self.assertEqual(set([0, 1]), set(sample_values.flatten()))
    # In this test we're just interested in verifying there isn't a crash
    # owing to mismatched types. b/30940152
    dist = tfd.Bernoulli(np.log([.2, .4]))
    self.assertAllEqual((1, 2), dist.sample(
        1, seed=tfp_test_util.test_seed()).shape.as_list())

  def testNotReparameterized(self):
    p = tf.constant([0.2, 0.6])
    _, grad_p = tfp.math.value_and_gradient(
        lambda x: tfd.Bernoulli(probs=x).sample(100), p)
    self.assertIsNone(grad_p)

  def testSampleDeterministicScalarVsVector(self):
    p = [0.2, 0.6]
    dist = tfd.Bernoulli(probs=p)
    n = 1000

    def _maybe_seed():
      if tf.executing_eagerly():
        tf.compat.v1.set_random_seed(42)
        return None
      return 42
    self.assertAllEqual(
        self.evaluate(dist.sample(n, _maybe_seed())),
        self.evaluate(dist.sample([n], _maybe_seed())))
    n = tf.compat.v1.placeholder_with_default(np.int32(1000), shape=None)
    sample1 = dist.sample(n, _maybe_seed())
    sample2 = dist.sample([n], _maybe_seed())
    sample1, sample2 = self.evaluate([sample1, sample2])
    self.assertAllEqual(sample1, sample2)

  def testMean(self):
    p = np.array([[0.2, 0.7], [0.5, 0.4]], dtype=np.float32)
    dist = tfd.Bernoulli(probs=p)
    self.assertAllEqual(self.evaluate(dist.mean()), p)

  def testVarianceAndStd(self):
    var = lambda p: p * (1. - p)
    p = [[0.2, 0.7], [0.5, 0.4]]
    dist = tfd.Bernoulli(probs=p)
    self.assertAllClose(
        self.evaluate(dist.variance()),
        np.array([[var(0.2), var(0.7)], [var(0.5), var(0.4)]],
                 dtype=np.float32))
    self.assertAllClose(
        self.evaluate(dist.stddev()),
        np.array([[np.sqrt(var(0.2)), np.sqrt(var(0.7))],
                  [np.sqrt(var(0.5)), np.sqrt(var(0.4))]],
                 dtype=np.float32))

  def testBernoulliBernoulliKL(self):
    batch_size = 6
    a_p = np.array([0.5] * batch_size, dtype=np.float32)
    b_p = np.array([0.4] * batch_size, dtype=np.float32)

    a = tfd.Bernoulli(probs=a_p)
    b = tfd.Bernoulli(probs=b_p)

    kl = tfd.kl_divergence(a, b)
    kl_val = self.evaluate(kl)

    kl_expected = (a_p * np.log(a_p / b_p) + (1. - a_p) * np.log(
        (1. - a_p) / (1. - b_p)))

    self.assertEqual(kl.shape, (batch_size,))
    self.assertAllClose(kl_val, kl_expected)

if __name__ == "__main__":
  tf.test.main()
