# Copyright 2018 DeepMind Technologies Limited.
#
# 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 mathematics_dataset.sample.ops."""

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

# Dependency imports
from absl.testing import absltest
from mathematics_dataset.sample import ops
from six.moves import range
import sympy


class OpsTest(absltest.TestCase):

  def testNeg(self):
    op = ops.Neg(2)
    self.assertEqual(str(op), '-2')
    self.assertEqual(op.sympy(), -2)

    op = ops.Add(ops.Neg(2), 3)
    self.assertEqual(str(op), '-2 + 3')
    self.assertEqual(op.sympy(), 1)

    op = ops.Add(3, ops.Neg(2))
    self.assertEqual(str(op), '3 - 2')
    self.assertEqual(op.sympy(), 1)

    op = ops.Add(ops.Add(ops.Neg(2), 5), 3)
    self.assertEqual(str(op), '-2 + 5 + 3')
    self.assertEqual(op.sympy(), 6)

    op = ops.Add(3, ops.Add(ops.Identity(ops.Neg(2)), 5))
    self.assertEqual(str(op), '3 - 2 + 5')
    self.assertEqual(op.sympy(), 6)

    op = ops.Add(3, ops.Add(2, ops.Neg(5)))
    self.assertEqual(str(op), '3 + 2 - 5')
    self.assertEqual(op.sympy(), 0)

  def testAdd(self):
    add = ops.Add()
    self.assertEqual(str(add), '0')
    self.assertEqual(add.sympy(), 0)

    add = ops.Add(2, 3)
    self.assertEqual(str(add), '2 + 3')
    self.assertEqual(add.sympy(), 5)

    add = ops.Add(ops.Add(1, 2), 3)
    self.assertEqual(str(add), '1 + 2 + 3')
    self.assertEqual(add.sympy(), 6)

  def testSub(self):
    sub = ops.Sub(2, 3)
    self.assertEqual(str(sub), '2 - 3')
    self.assertEqual(sub.sympy(), -1)

    sub = ops.Sub(ops.Sub(1, 2), 3)
    self.assertEqual(str(sub), '1 - 2 - 3')
    self.assertEqual(sub.sympy(), -4)

    sub = ops.Sub(1, ops.Sub(2, 3))
    self.assertEqual(str(sub), '1 - (2 - 3)')
    self.assertEqual(sub.sympy(), 2)

    sub = ops.Sub(ops.Neg(1), 2)
    self.assertEqual(str(sub), '-1 - 2')
    self.assertEqual(sub.sympy(), -3)

  def testMul(self):
    mul = ops.Mul()
    self.assertEqual(str(mul), '1')
    self.assertEqual(mul.sympy(), 1)

    mul = ops.Mul(2, 3)
    self.assertEqual(str(mul), '2*3')
    self.assertEqual(mul.sympy(), 6)

    mul = ops.Mul(ops.Identity(ops.Constant(-2)), 3)
    self.assertEqual(str(mul), '-2*3')
    self.assertEqual(mul.sympy(), -6)

    mul = ops.Mul(ops.Add(1, 2), 3)
    self.assertEqual(str(mul), '(1 + 2)*3')
    self.assertEqual(mul.sympy(), 9)

    mul = ops.Mul(ops.Mul(2, 3), 5)
    self.assertEqual(str(mul), '2*3*5')
    self.assertEqual(mul.sympy(), 30)

    # TODO(b/124038946): reconsider how we want brackets in these cases:
#     mul = ops.Mul(ops.Div(2, 3), 5)
#     self.assertEqual(str(mul), '(2/3)*5')
#     self.assertEqual(mul.sympy(), sympy.Rational(10, 3))
#
#     mul = ops.Mul(sympy.Rational(2, 3), 5)
#     self.assertEqual(str(mul), '(2/3)*5')
#     self.assertEqual(mul.sympy(), sympy.Rational(10, 3))

  def testDiv(self):
    div = ops.Div(2, 3)
    self.assertEqual(str(div), '2/3')
    self.assertEqual(div.sympy(), sympy.Rational(2, 3))

    div = ops.Div(2, sympy.Rational(4, 5))
    self.assertEqual(str(div), '2/(4/5)')
    self.assertEqual(div.sympy(), sympy.Rational(5, 2))

    div = ops.Div(1, ops.Div(2, 3))
    self.assertEqual(str(div), '1/(2/3)')
    self.assertEqual(div.sympy(), sympy.Rational(3, 2))

    div = ops.Div(ops.Div(2, 3), 4)
    self.assertEqual(str(div), '(2/3)/4')
    self.assertEqual(div.sympy(), sympy.Rational(1, 6))

    div = ops.Div(2, ops.Mul(3, 4))
    self.assertEqual(str(div), '2/(3*4)')

    div = ops.Div(2, sympy.Function('f')(sympy.Symbol('x')))
    self.assertEqual(str(div), '2/f(x)')

  def testPow(self):
    pow_ = ops.Pow(2, 3)
    self.assertEqual(str(pow_), '2**3')
    self.assertEqual(pow_.sympy(), 8)

    pow_ = ops.Pow(4, sympy.Rational(1, 2))
    self.assertEqual(str(pow_), '4**(1/2)')
    self.assertEqual(pow_.sympy(), 2)

    pow_ = ops.Pow(sympy.Rational(1, 2), 3)
    self.assertEqual(str(pow_), '(1/2)**3')
    self.assertEqual(pow_.sympy(), 1/8)

    pow_ = ops.Pow(3, ops.Pow(2, 1))
    self.assertEqual(str(pow_), '3**(2**1)')
    self.assertEqual(pow_.sympy(), 9)

    pow_ = ops.Pow(ops.Pow(2, 3), 4)
    self.assertEqual(str(pow_), '(2**3)**4')
    self.assertEqual(pow_.sympy(), 4096)

    pow_ = ops.Pow(-5, 2)
    self.assertEqual(str(pow_), '(-5)**2')
    self.assertEqual(pow_.sympy(), 25)

  def testEq(self):
    op = ops.Eq(ops.Add(2, 3), 4)
    self.assertEqual(str(op), '2 + 3 = 4')
    self.assertEqual(op.sympy(), False)

  def testDescendants(self):
    constants = [ops.Constant(i) for i in range(6)]

    # (1 + 2*3**4) / 5 - 6
    expression = ops.Sub(
        ops.Div(
            ops.Add(
                constants[0],
                ops.Mul(
                    constants[1],
                    ops.Pow(
                        constants[2],
                        constants[3]))),
            constants[4]),
        constants[5])
    descendants = expression.descendants()
    descendants = ops._flatten(descendants)

    for constant in constants:
      self.assertIn(constant, descendants)
      self.assertEqual(descendants.count(constant), 1)

    # Also test top-level.
    self.assertEqual(constants[0].descendants(), [constants[0]])

    # Also general structure.
    constant = ops.Constant(3)
    expression = ops.Neg(constant)
    self.assertEqual(set(expression.descendants()), set([constant, expression]))

  def testNumberConstants(self):
    constant = ops.Constant(3)
    expression = ops.Neg(constant)
    constants = ops.number_constants([expression])
    self.assertEqual(constants, [constant])


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