# Copyright 2016 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 tensorflow.python.ops.special_math_ops."""

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

import numpy as np

from tensorflow.python.client import session
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.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging

class LBetaTest(test.TestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_one_dimensional_arg(self):
    # Should evaluate to 1 and 1/2.
    x_one = [1, 1.]
    x_one_half = [2, 1.]
    with self.session(use_gpu=True):
      self.assertAllClose(
          1, self.evaluate(math_ops.exp(special_math_ops.lbeta(x_one))))
      self.assertAllClose(
          0.5, self.evaluate(math_ops.exp(special_math_ops.lbeta(x_one_half))))
      self.assertEqual([], special_math_ops.lbeta(x_one).get_shape())

  @test_util.run_deprecated_v1
  def test_one_dimensional_arg_dynamic(self):
    # Should evaluate to 1 and 1/2.
    x_one = [1, 1.]
    x_one_half = [2, 1.]
    with self.session(use_gpu=True):
      ph = array_ops.placeholder(dtypes.float32)
      beta_ph = math_ops.exp(special_math_ops.lbeta(ph))
      self.assertAllClose(1, beta_ph.eval(feed_dict={ph: x_one}))
      self.assertAllClose(0.5,
                          beta_ph.eval(feed_dict={ph: x_one_half}))

  @test_util.run_deprecated_v1
  def test_four_dimensional_arg_with_partial_shape_dynamic(self):
    x_ = np.ones((3, 2, 3, 4))
    # Gamma(1) = 0! = 1
    # Gamma(1 + 1 + 1 + 1) = Gamma(4) = 3! = 6
    # ==> Beta([1, 1, 1, 1])
    #     = Gamma(1) * Gamma(1) * Gamma(1) * Gamma(1) / Gamma(1 + 1 + 1 + 1)
    #     = 1 / 6
    expected_beta_x = 1 / 6 * np.ones((3, 2, 3))
    with self.session(use_gpu=True):
      x_ph = array_ops.placeholder(dtypes.float32, [3, 2, 3, None])
      beta_ph = math_ops.exp(special_math_ops.lbeta(x_ph))
      self.assertAllClose(expected_beta_x,
                          beta_ph.eval(feed_dict={x_ph: x_}))

  @test_util.run_in_graph_and_eager_modes
  def test_two_dimensional_arg(self):
    # Should evaluate to 1/2.
    x_one_half = [[2, 1.], [2, 1.]]
    with self.session(use_gpu=True):
      self.assertAllClose(
          [0.5, 0.5],
          self.evaluate(math_ops.exp(special_math_ops.lbeta(x_one_half))))
      self.assertEqual((2,), special_math_ops.lbeta(x_one_half).get_shape())

  @test_util.run_deprecated_v1
  def test_two_dimensional_arg_dynamic(self):
    # Should evaluate to 1/2.
    x_one_half = [[2, 1.], [2, 1.]]
    with self.session(use_gpu=True):
      ph = array_ops.placeholder(dtypes.float32)
      beta_ph = math_ops.exp(special_math_ops.lbeta(ph))
      self.assertAllClose([0.5, 0.5],
                          beta_ph.eval(feed_dict={ph: x_one_half}))

  @test_util.run_in_graph_and_eager_modes
  def test_two_dimensional_proper_shape(self):
    # Should evaluate to 1/2.
    x_one_half = [[2, 1.], [2, 1.]]
    with self.session(use_gpu=True):
      self.assertAllClose(
          [0.5, 0.5],
          self.evaluate(math_ops.exp(special_math_ops.lbeta(x_one_half))))
      self.assertEqual(
          (2,),
          self.evaluate(array_ops.shape(special_math_ops.lbeta(x_one_half))))
      self.assertEqual(
          tensor_shape.TensorShape([2]),
          special_math_ops.lbeta(x_one_half).get_shape())

  @test_util.run_in_graph_and_eager_modes
  def test_complicated_shape(self):
    with self.session(use_gpu=True):
      x = ops.convert_to_tensor(np.random.rand(3, 2, 2))
      self.assertAllEqual(
          (3, 2), self.evaluate(array_ops.shape(special_math_ops.lbeta(x))))
      self.assertEqual(
          tensor_shape.TensorShape([3, 2]),
          special_math_ops.lbeta(x).get_shape())

  @test_util.run_in_graph_and_eager_modes
  def test_length_1_last_dimension_results_in_one(self):
    # If there is only one coefficient, the formula still works, and we get one
    # as the answer, always.
    x_a = [5.5]
    x_b = [0.1]
    with self.session(use_gpu=True):
      self.assertAllClose(
          1,
          self.evaluate(math_ops.exp(special_math_ops.lbeta(x_a))),
          rtol=3e-6)
      self.assertAllClose(
          1, self.evaluate(math_ops.exp(special_math_ops.lbeta(x_b))))
      self.assertEqual((), special_math_ops.lbeta(x_a).get_shape())

  @test_util.run_in_graph_and_eager_modes
  def test_empty_rank1_returns_negative_infinity(self):
    with self.session(use_gpu=True):
      x = constant_op.constant([], shape=[0])
      lbeta_x = special_math_ops.lbeta(x)
      expected_result = constant_op.constant(-np.inf, shape=())

      self.assertAllEqual(self.evaluate(expected_result),
                          self.evaluate(lbeta_x))
      self.assertEqual(expected_result.get_shape(), lbeta_x.get_shape())

  @test_util.run_in_graph_and_eager_modes
  def test_empty_rank2_with_zero_last_dim_returns_negative_infinity(self):
    with self.session(use_gpu=True):
      event_size = 0
      for batch_size in [0, 1, 2]:
        x = constant_op.constant([], shape=[batch_size, event_size])
        lbeta_x = special_math_ops.lbeta(x)
        expected_result = constant_op.constant(-np.inf, shape=[batch_size])

        self.assertAllEqual(self.evaluate(expected_result),
                            self.evaluate(lbeta_x))
        self.assertEqual(expected_result.get_shape(), lbeta_x.get_shape())

  @test_util.run_in_graph_and_eager_modes
  def test_empty_rank2_with_zero_batch_dim_returns_empty(self):
    with self.session(use_gpu=True):
      batch_size = 0
      for event_size in [0, 1, 2]:
        x = constant_op.constant([], shape=[batch_size, event_size])
        lbeta_x = special_math_ops.lbeta(x)

        expected_result = constant_op.constant([], shape=[batch_size])

        self.assertAllEqual(self.evaluate(expected_result),
                            self.evaluate(lbeta_x))
        self.assertEqual(expected_result.get_shape(), lbeta_x.get_shape())


class BesselTest(test.TestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_bessel_i0(self):
    x_single = np.arange(-3, 3).reshape(1, 3, 2).astype(np.float32)
    x_double = np.arange(-3, 3).reshape(1, 3, 2).astype(np.float64)
    try:
      from scipy import special  # pylint: disable=g-import-not-at-top
      self.assertAllClose(special.i0(x_single),
                          self.evaluate(special_math_ops.bessel_i0(x_single)))
      self.assertAllClose(special.i0(x_double),
                          self.evaluate(special_math_ops.bessel_i0(x_double)))
    except ImportError as e:
      tf_logging.warn('Cannot test special functions: %s' % str(e))

  @test_util.run_in_graph_and_eager_modes
  def test_bessel_i1(self):
    x_single = np.arange(-3, 3).reshape(1, 3, 2).astype(np.float32)
    x_double = np.arange(-3, 3).reshape(1, 3, 2).astype(np.float64)
    try:
      from scipy import special  # pylint: disable=g-import-not-at-top
      self.assertAllClose(special.i1(x_single),
                          self.evaluate(special_math_ops.bessel_i1(x_single)))
      self.assertAllClose(special.i1(x_double),
                          self.evaluate(special_math_ops.bessel_i1(x_double)))
    except ImportError as e:
      tf_logging.warn('Cannot test special functions: %s' % str(e))


class EinsumTest(test.TestCase):

  simple_cases = [
      'ij,jk->ik', 'ijk,jklm->il', 'ij,jk,kl->il', 'ijk->i', 'ijk->kji',
      'ji,kj->ik', 'ikl,kji->kl', 'klj,lki->ij', 'ijk,ilj->kli',
      'kij,mkb->ijmb', 'ijk,ijl,ikl->i', 'i,ijk,j->k', 'ij,ij,jk,kl->il',
      'ij,kj,il,jm->ml', 'a,ab,abc->abc', 'a,b,ab->ab', 'ab,ab,c->',
      'ab,ab,c->c', 'ab,ab,cd,cd->', 'ab,ab,cd,cd->ac', 'ab,ab,cd,cd->cd',
      'ab,ab,cd,cd,ef,ef->', 'ab,cd,ef->abcdef', 'ab,cd,ef->acdf',
      'ab,cd,de->abcde', 'ab,cd,de->be', 'ab,bcd,cd->abcd', 'ab,bcd,cd->abd',
      'eb,cb,fb->cef', 'abcd,ad', 'bd,db,eac->ace', 'ba,ac,da->bcd', 'ab,ab',
      'ab,ba', 'abc,abc', 'abc,bac', 'abc,cba', 'dba,ead,cad->bce',
      'aef,fbc,dca->bde', 'iJ,Jk->ik', 'iJ,Ki->JK', 'iJk,Jklm->Jk',
      'ij, jk, kl -> il', 'a, ab, abc -> abc', 'ab, ab, cd, cd, ef, ef -> ',
      'abc, bac', 'iJ, Ki -> JK', 'iJk, Jklm -> Jk', 'ii', 'ijji'
  ]

  long_cases = [
      'bca,cdb,dbf,afc->',
      'efc,dbc,acf,fd->abe',
      'ea,fb,gc,hd,abcd->efgh',
      'ea,fb,abcd,gc,hd->efgh',
      'abhe,hidj,jgba,hiab,gab',
      'efc, dbc, acf, fd -> abe',
      'abhe, hidj, jgba, hiab, gab',
  ]

  invalid_cases = [
      # bad formats
      '',
      'ijk ijk',
      'ij.jk->ik',
      '...ij...,jk...->ik...',
      '...ij,jk...->...ik...',
      'i$...',
      'ij...,jk...->ik.',
      'ij.,jk...->ik...',
      'ij,k ->kji',
      'ij,k-> kji',

      # axis in output that does not exist
      'ij,jk->im',

      # incorrect number of dimensions
      'ij,jkl->kl',

      # this is allowed in numpy but not implemented here yet
      'iij,jk',
      'a...a->a...',
  ]

  ellipsis_cases = [
      ['...ij,...jk->...ik', 'abij,abjk->abik'],  # batch mat-mat multiply
      ['...ij,...jk->...ik', 'abij,bjk->abik'],
      ['...ij,...jk->...ik', 'abij,jk->abik'],
      ['...ij,...jk->...ik', 'ij,jk->ik'],
      ['...ij,...jk->...ik', 'abij,abjk->'],
      ['...ij,...jk->', 'bij,bjk->'],
      ['...ij,...jk', 'bij,bjk'],
      ['...XY,...YZ', 'bXY,bYZ'],
      ['...ik,...i->...i', 'abik,abi->abi'],      # batch mat-vec multiply
      ['...ij->...ji', 'aij->aji'],               # batch matrix transpose
      ['...ij->...', 'aij->a'],                   # batch sum
      ['...,...->...', 'ij,ij->ij'],              # hadamard product
      ['...i,...j->...ij', 'abi,abj->abij'],      # outer product
      ['i...,j...->...ij', 'iab,jab->abij'],      #
      ['i...,j...->ij...', 'iab,jab->ijab'],      #
      ['...ik,...jkl,...il->...ij', 'aik,ajkl,ail->aij'],
      ['...', 'abc'],                             # copy
      ['...ca,cd...,d...f,afc->', 'bca,cdb,dbf,afc->'],
      # Tests based on https://github.com/dask/dask/pull/3412/files
      ['ab...,bc...->ac...', 'abxy,bcxy->acxy'],
      ['a...a', 'aba'],
      ['a...a', 'aa'],
      ['abc...->cba...', 'abcxy->cbaxy'],
      ['...ab->...a', 'xyab->xya'],
      ['...abc,...abcd->...d', 'xyabc,xabcd->xyd'],
      ['ab...,b->ab...', 'abxy,b->abxy'],
  ]

  dim_mismatch_cases = [('ijk,jkl->il', [(2, 3, 4), (3, 5, 6)]),
                        ('...ij,...jk->...ik', [(3), (3, 4)]),
                        (('...abcdefghijklmnopqrstuvwxyz,'
                          '...ABCDEFGHIJKLMNOPQRSTUVWXYZ'),
                         [tuple([1] * 27), tuple([1] * 27)])]

  def test_simple(self):
    for case in self.simple_cases:
      self.run_test(case)

  def test_long(self):
    for case in self.long_cases:
      self.run_test(case)

  def test_ellipsis(self):
    for (axes, expanded_axes) in self.ellipsis_cases:
      self.run_test(axes, expanded_axes)

  @test_util.run_deprecated_v1
  def test_invalid(self):
    for axes in self.invalid_cases:
      inputs = [
          array_ops.placeholder(dtypes.float32, shape=(3, 4)),
          array_ops.placeholder(dtypes.float32, shape=(3, 4)),
      ]
      with self.assertRaises(ValueError):
        _ = special_math_ops.einsum(axes, *inputs)

  @test_util.run_deprecated_v1
  def test_invalid_keyword_arguments(self):
    m0 = array_ops.placeholder(dtypes.int32, shape=(1, None))
    m1 = array_ops.placeholder(dtypes.int32, shape=(None, 1))
    with self.assertRaisesRegexp(
        TypeError,
        'invalid keyword arguments for this function: invalid1, invalid2'):
      _ = special_math_ops.einsum(
          'ij,jk->ik',
          m0,
          m1,
          name='name',
          invalid1='value1',
          invalid2='value2')

  @test_util.run_deprecated_v1
  def test_repeated_axis_single_input(self):
    x = array_ops.placeholder(dtypes.float32, shape=[2, 2])
    with self.assertRaises(ValueError):
      _ = special_math_ops.einsum('ii->', x)

  @test_util.run_deprecated_v1
  def test_dim_mismatch(self):
    for axes, input_shapes in self.dim_mismatch_cases:
      inputs = [
          array_ops.placeholder(dtypes.float32, shape=shape)
          for shape in input_shapes
      ]
      with self.assertRaises(ValueError):
        _ = special_math_ops.einsum(axes, *inputs)

  def run_test(self, axes, expanded_axes=None):
    expanded_axes = expanded_axes if expanded_axes is not None else axes
    all_axes = {ax: np.random.randint(4, 12)
                for ax in expanded_axes if ax.isalpha()}

    input_vals = []
    input_axes, _, _ = axes.partition('->')

    for idx in input_axes.split(','):
      shape = [all_axes[ax] for ax in idx if ax.isalpha()]
      input_vals.append(np.random.random(shape))

    input_tensors = [constant_op.constant(val) for val in input_vals]
    output_tensor = special_math_ops.einsum(axes, *input_tensors)

    with self.session(use_gpu=True):
      output_value = self.evaluate(output_tensor)

    correct_value = 0
    if axes == 'ijji':
      output = math_ops.trace(*input_tensors)
      correct_value = self.evaluate(output)
    else:
      correct_value = np.einsum(axes, *input_vals)
    err = np.abs(correct_value - output_value).max()
    self.assertLess(err, 1e-8)

  def test_input_is_placeholder(self):
    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(1, None))
      m1 = array_ops.placeholder(dtypes.int32, shape=(None, 1))
      out = special_math_ops.einsum('ij,jk->ik', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[1, 2, 3]],
            m1: [[2], [1], [1]],
        }
        self.assertAllClose([[7]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, 3))
      m1 = array_ops.placeholder(dtypes.int32, shape=(3,))
      out = special_math_ops.einsum('ij,j->i', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[1, 2, 3]],
            m1: [2, 1, 1],
        }
        self.assertAllClose([7], sess.run(out, feed_dict=feed_dict))

    # Tests for placeholders which have two or more None values
    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, None, 2))
      m1 = array_ops.placeholder(dtypes.int32, shape=(2, 1))
      out = special_math_ops.einsum('ijk,kl->ijl', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[[1, 2]]],
            m1: [[3], [2]],
        }
        self.assertAllClose([[[7]]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(2, 1))
      m1 = array_ops.placeholder(dtypes.int32, shape=(None, None, 2))
      out = special_math_ops.einsum('kl,ijk->ijl', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[3], [2]],
            m1: [[[1, 2]]],
        }
        self.assertAllClose([[[7]]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, None, 2))
      m1 = array_ops.placeholder(dtypes.int32, shape=(2,))
      out = special_math_ops.einsum('ijk,k->ij', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[[1, 2]]],
            m1: [3, 2],
        }
        self.assertAllClose([[7]], sess.run(out, feed_dict=feed_dict))

    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.int32, shape=(None, 2, None, 2))
      m1 = array_ops.placeholder(dtypes.int32, shape=(None, 2))
      out = special_math_ops.einsum('ijkl,ij->ikl', m0, m1)
      with session.Session() as sess:
        feed_dict = {
            m0: [[[[1, 2]], [[2, 1]]]],
            m1: [[3, 2]],
        }
        self.assertAllClose([[[7, 8]]], sess.run(out, feed_dict=feed_dict))

  @test_util.run_in_graph_and_eager_modes
  def test_multiple_ellipses(self):
    m0 = array_ops.placeholder_with_default([[[[1, 2]], [[2, 1]]]],
                                            shape=(None, 2, None, 2))
    m1 = array_ops.placeholder_with_default([[3, 2]], shape=(None, 2))
    out = special_math_ops.einsum('...jkl,...j->...kl', m0, m1)
    self.assertAllClose([[[7, 8]]], self.evaluate(out))

  def test_ellipses_with_unknown_input_dim(self):
    with ops.Graph().as_default():
      m0 = array_ops.placeholder(dtypes.float32)
      m1 = array_ops.placeholder_with_default([[3, 2]], shape=(None, 2))
      with self.assertRaises(ValueError):
        _ = special_math_ops.einsum('...jkl,...j->...kl', m0, m1)

if __name__ == '__main__':
  test.main()
