# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# 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 initializers."""

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

import importlib
import math

import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops.distributions import kullback_leibler
from tensorflow.python.ops.distributions import normal as normal_lib
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging


def try_import(name):  # pylint: disable=invalid-name
  module = None
  try:
    module = importlib.import_module(name)
  except ImportError as e:
    tf_logging.warning("Could not import %s: %s" % (name, str(e)))
  return module

stats = try_import("scipy.stats")


class NormalTest(test.TestCase):

  def setUp(self):
    self._rng = np.random.RandomState(123)

  def assertAllFinite(self, tensor):
    is_finite = np.isfinite(tensor.eval())
    all_true = np.ones_like(is_finite, dtype=np.bool)
    self.assertAllEqual(all_true, is_finite)

  def _testParamShapes(self, sample_shape, expected):
    with self.test_session():
      param_shapes = normal_lib.Normal.param_shapes(sample_shape)
      mu_shape, sigma_shape = param_shapes["loc"], param_shapes["scale"]
      self.assertAllEqual(expected, mu_shape.eval())
      self.assertAllEqual(expected, sigma_shape.eval())
      mu = array_ops.zeros(mu_shape)
      sigma = array_ops.ones(sigma_shape)
      self.assertAllEqual(
          expected,
          array_ops.shape(normal_lib.Normal(mu, sigma).sample()).eval())

  def _testParamStaticShapes(self, sample_shape, expected):
    param_shapes = normal_lib.Normal.param_static_shapes(sample_shape)
    mu_shape, sigma_shape = param_shapes["loc"], param_shapes["scale"]
    self.assertEqual(expected, mu_shape)
    self.assertEqual(expected, sigma_shape)

  def testParamShapes(self):
    sample_shape = [10, 3, 4]
    self._testParamShapes(sample_shape, sample_shape)
    self._testParamShapes(constant_op.constant(sample_shape), sample_shape)

  def testParamStaticShapes(self):
    sample_shape = [10, 3, 4]
    self._testParamStaticShapes(sample_shape, sample_shape)
    self._testParamStaticShapes(
        tensor_shape.TensorShape(sample_shape), sample_shape)

  def testNormalWithSoftplusScale(self):
    with self.test_session():
      mu = array_ops.zeros((10, 3))
      rho = array_ops.ones((10, 3)) * -2.
      normal = normal_lib.NormalWithSoftplusScale(loc=mu, scale=rho)
      self.assertAllEqual(mu.eval(), normal.loc.eval())
      self.assertAllEqual(nn_ops.softplus(rho).eval(), normal.scale.eval())

  def testNormalLogPDF(self):
    with self.test_session():
      batch_size = 6
      mu = constant_op.constant([3.0] * batch_size)
      sigma = constant_op.constant([math.sqrt(10.0)] * batch_size)
      x = np.array([-2.5, 2.5, 4.0, 0.0, -1.0, 2.0], dtype=np.float32)
      normal = normal_lib.Normal(loc=mu, scale=sigma)

      log_pdf = normal.log_prob(x)
      self.assertAllEqual(normal.batch_shape_tensor().eval(),
                          log_pdf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(),
                          log_pdf.eval().shape)
      self.assertAllEqual(normal.batch_shape, log_pdf.get_shape())
      self.assertAllEqual(normal.batch_shape, log_pdf.eval().shape)

      pdf = normal.prob(x)
      self.assertAllEqual(normal.batch_shape_tensor().eval(), pdf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(), pdf.eval().shape)
      self.assertAllEqual(normal.batch_shape, pdf.get_shape())
      self.assertAllEqual(normal.batch_shape, pdf.eval().shape)

      if not stats:
        return
      expected_log_pdf = stats.norm(mu.eval(), sigma.eval()).logpdf(x)
      self.assertAllClose(expected_log_pdf, log_pdf.eval())
      self.assertAllClose(np.exp(expected_log_pdf), pdf.eval())

  def testNormalLogPDFMultidimensional(self):
    with self.test_session():
      batch_size = 6
      mu = constant_op.constant([[3.0, -3.0]] * batch_size)
      sigma = constant_op.constant([[math.sqrt(10.0), math.sqrt(15.0)]] *
                                   batch_size)
      x = np.array([[-2.5, 2.5, 4.0, 0.0, -1.0, 2.0]], dtype=np.float32).T
      normal = normal_lib.Normal(loc=mu, scale=sigma)

      log_pdf = normal.log_prob(x)
      log_pdf_values = log_pdf.eval()
      self.assertEqual(log_pdf.get_shape(), (6, 2))
      self.assertAllEqual(normal.batch_shape_tensor().eval(),
                          log_pdf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(),
                          log_pdf.eval().shape)
      self.assertAllEqual(normal.batch_shape, log_pdf.get_shape())
      self.assertAllEqual(normal.batch_shape, log_pdf.eval().shape)

      pdf = normal.prob(x)
      pdf_values = pdf.eval()
      self.assertEqual(pdf.get_shape(), (6, 2))
      self.assertAllEqual(normal.batch_shape_tensor().eval(), pdf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(), pdf_values.shape)
      self.assertAllEqual(normal.batch_shape, pdf.get_shape())
      self.assertAllEqual(normal.batch_shape, pdf_values.shape)

      if not stats:
        return
      expected_log_pdf = stats.norm(mu.eval(), sigma.eval()).logpdf(x)
      self.assertAllClose(expected_log_pdf, log_pdf_values)
      self.assertAllClose(np.exp(expected_log_pdf), pdf_values)

  def testNormalCDF(self):
    with self.test_session():
      batch_size = 50
      mu = self._rng.randn(batch_size)
      sigma = self._rng.rand(batch_size) + 1.0
      x = np.linspace(-8.0, 8.0, batch_size).astype(np.float64)

      normal = normal_lib.Normal(loc=mu, scale=sigma)
      cdf = normal.cdf(x)
      self.assertAllEqual(normal.batch_shape_tensor().eval(), cdf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(), cdf.eval().shape)
      self.assertAllEqual(normal.batch_shape, cdf.get_shape())
      self.assertAllEqual(normal.batch_shape, cdf.eval().shape)
      if not stats:
        return
      expected_cdf = stats.norm(mu, sigma).cdf(x)
      self.assertAllClose(expected_cdf, cdf.eval(), atol=0)

  def testNormalSurvivalFunction(self):
    with self.test_session():
      batch_size = 50
      mu = self._rng.randn(batch_size)
      sigma = self._rng.rand(batch_size) + 1.0
      x = np.linspace(-8.0, 8.0, batch_size).astype(np.float64)

      normal = normal_lib.Normal(loc=mu, scale=sigma)

      sf = normal.survival_function(x)
      self.assertAllEqual(normal.batch_shape_tensor().eval(), sf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(), sf.eval().shape)
      self.assertAllEqual(normal.batch_shape, sf.get_shape())
      self.assertAllEqual(normal.batch_shape, sf.eval().shape)
      if not stats:
        return
      expected_sf = stats.norm(mu, sigma).sf(x)
      self.assertAllClose(expected_sf, sf.eval(), atol=0)

  def testNormalLogCDF(self):
    with self.test_session():
      batch_size = 50
      mu = self._rng.randn(batch_size)
      sigma = self._rng.rand(batch_size) + 1.0
      x = np.linspace(-100.0, 10.0, batch_size).astype(np.float64)

      normal = normal_lib.Normal(loc=mu, scale=sigma)

      cdf = normal.log_cdf(x)
      self.assertAllEqual(normal.batch_shape_tensor().eval(), cdf.get_shape())
      self.assertAllEqual(normal.batch_shape_tensor().eval(), cdf.eval().shape)
      self.assertAllEqual(normal.batch_shape, cdf.get_shape())
      self.assertAllEqual(normal.batch_shape, cdf.eval().shape)

      if not stats:
        return
      expected_cdf = stats.norm(mu, sigma).logcdf(x)
      self.assertAllClose(expected_cdf, cdf.eval(), atol=0, rtol=1e-5)

  def testFiniteGradientAtDifficultPoints(self):
    for dtype in [np.float32, np.float64]:
      g = ops.Graph()