# 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.polynomials."""

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

import random

# Dependency imports
from absl.testing import parameterized
from mathematics_dataset.sample import polynomials
import numpy as np
from six.moves import range
import sympy
import tensorflow as tf


class ExpressionWithValueTest(tf.test.TestCase, parameterized.TestCase):

  def testSplitValueEqually(self):
    split = polynomials._split_value_equally(3, 2)
    self.assertEqual(split, [1, 2])
    split = polynomials._split_value_equally(sympy.sympify('3/4'), 2)
    self.assertEqual(split, [sympy.sympify('1/4'), sympy.sympify('1/2')])

  def testIntegersWithSum(self):
    value = 13
    count = 10
    terms = polynomials.integers_with_sum(value=value, count=count, entropy=4.0)
    self.assertLen(terms, count)
    self.assertEqual(sum(terms), value)

  def testMonomial(self):
    x, y = sympy.symbols('x y')
    self.assertEqual(str(polynomials.monomial(1, [x, y], [2, 3])), 'x**2*y**3')
    # TODO(b/124038530): how handle rational coefficients; are they even used?
    # self.assertEqual(
    #     str(polynomials.monomial(sympy.Rational(2, 3), [x], [1])), '2*x/3')
    # self.assertEqual(
    #     str(polynomials.monomial(sympy.Rational(1, 3), [x], [1])), 'x/3')
    self.assertEqual(str(polynomials.monomial(x, [y], [4])), 'x*y**4')

  def testExpandCoefficients(self):
    for _ in range(10):
      num_variables = np.random.randint(1, 4)
      degrees = np.random.randint(0, 4, [num_variables])
      coefficients = np.random.randint(-3, 3, degrees + 1)
      entropy = np.random.uniform(0, 10)
      expanded = polynomials.expand_coefficients(coefficients, entropy)
      collapsed = np.vectorize(sum)(expanded)
      self.assertAllEqual(coefficients, collapsed)

  def testCoefficientsToPolynomial(self):
    coeffs = [3, 2, 1]
    x = sympy.Symbol('x')
    polynomial = polynomials.coefficients_to_polynomial(coeffs, [x])
    polynomial = sympy.sympify(polynomial)
    self.assertEqual(polynomial, x*x + 2*x + 3)

  def testUnivariate(self):
    # Test generation for: x**2 + 2*x + 1
    x = sympy.Symbol('x')
    coeffs = [1, 2, 3]
    for _ in range(10):
      expanded = polynomials.expand_coefficients(coeffs, 5.0)
      polynomial = polynomials.coefficients_to_polynomial(expanded, [x])
      sympified = sympy.sympify(polynomial)
      self.assertEqual(sympified, 1 + 2*x + 3*x*x)

  def testMultivariate(self):
    # Test generation for: x**2 + 2*x*y + 3*y**2 - x + 5
    x, y = sympy.symbols('x y')
    coeffs = [[5, 0, 3], [-1, 2, 0], [1, 0, 0]]
    for _ in range(10):
      expanded = polynomials.expand_coefficients(coeffs, 5.0, length=10)
      polynomial = polynomials.coefficients_to_polynomial(expanded, [x, y])
      sympified = sympy.sympify(polynomial)
      self.assertEqual(sympified, x*x + 2*x*y + 3*y*y - x + 5)

  def testAddCoefficients(self):
    # Add x**2 + 2*y and 3*x + 4*y**3.
    coeffs1 = [[0, 2], [0, 0], [1, 0]]
    coeffs2 = [[0, 0, 0, 4], [3, 0, 0, 0]]
    target = [[0, 2, 0, 4], [3, 0, 0, 0], [1, 0, 0, 0]]
    actual = polynomials.add_coefficients(coeffs1, coeffs2)
    self.assertAllEqual(target, actual)

  def testCoefficientsLinearSplit(self):
    for degree in range(3):
      for ndims in range(3):
        for _ in range(10):
          coefficients = np.random.randint(-5, 5, [degree + 1] * ndims)
          entropy = random.uniform(1, 4)
          c1, c2, coeffs1, coeffs2 = polynomials.coefficients_linear_split(
              coefficients, entropy)
          c1 = int(c1)
          c2 = int(c2)
          coeffs1 = np.asarray(coeffs1, dtype=np.int32)
          coeffs2 = np.asarray(coeffs2, dtype=np.int32)
          sum_ = c1 * coeffs1 + c2 * coeffs2
          self.assertAllEqual(sum_, coefficients)

  def testSampleWithBrackets(self):
    x, y = sympy.symbols('x y')
    for _ in range(100):
      degrees = np.random.randint(1, 4, [2])
      entropy = random.uniform(0, 4)
      polynomial = polynomials.sample_with_brackets(
          variables=[x, y], degrees=degrees, entropy=entropy)
      self.assertIn('(', str(polynomial))
      poly = sympy.poly(sympy.sympify(polynomial).expand())
      self.assertEqual(poly.degree(x), degrees[0])
      self.assertEqual(poly.degree(y), degrees[1])

  def testTrim(self):
    self.assertAllEqual(polynomials.trim([1]), [1])
    self.assertAllEqual(polynomials.trim([1, 0]), [1])
    self.assertAllEqual(polynomials.trim([0, 1]), [0, 1])
    self.assertAllEqual(polynomials.trim([0]), [])
    self.assertAllEqual(polynomials.trim([0, 0]), [])

  def testDifferentiate_univariate(self):
    coeffs = [5, 3, 2]
    expected = [3, 4]
    actual = polynomials.differentiate(coeffs, 0)
    self.assertAllEqual(expected, actual)

  def testDifferentiate_multivariate(self):
    coeffs = [[0, 3, 1], [5, 0, 0], [0, 2, 0]]
    expected = [[5, 0], [0, 4]]
    actual = polynomials.differentiate(coeffs, 0)
    self.assertAllEqual(expected, actual)

  def testIntegrate_univariate(self):
    coeffs = [5, 3, 2]
    expected = [0, 5, sympy.Rational(3, 2), sympy.Rational(2, 3)]
    actual = polynomials.integrate(coeffs, 0)
    self.assertAllEqual(expected, actual)

  def testIntegrate_multivariate(self):
    coeffs = [[0, 1], [1, 0]]
    expected = [[0, 0, sympy.Rational(1, 2)], [0, 1, 0]]
    actual = polynomials.integrate(coeffs, 1)
    self.assertAllEqual(expected, actual)


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