# coding=utf-8
# Copyright 2019 The SEED Authors
# 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 seed_rl.grpc.python.ops."""

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

import collections
from concurrent import futures
import threading
import time
import uuid

from absl.testing import parameterized
import numpy as np
from seed_rl.grpc.python import ops
from six.moves import range
import tensorflow as tf

Some = collections.namedtuple('Some', 'a b')


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

  def get_unix_address(self):
    return 'unix:/tmp/%s' % uuid.uuid4()

  def test_simple(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(43, client.foo(42))
    server.shutdown()



  def test_simple_two_calls(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(43, client.foo(42))
    self.assertAllEqual(44, client.foo(43))
    server.shutdown()

  def test_empty_input(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[])
    def foo():
      return 42

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(42, client.foo())
    server.shutdown()

  def test_empty_output(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):  
      return []

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual([], client.foo(42))
    server.shutdown()

  def test_no_output(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):  
      pass

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertIsNone(client.foo(42))
    server.shutdown()

  def test_large_tensor(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    t = tf.fill([10, 1024, 1024], 1)  # 40MB

    @tf.function(input_signature=[tf.TensorSpec([] + list(t.shape), tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(t + 1, client.foo(t))
    server.shutdown()

  def test_create_variable(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    state = [None]

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      if state[0] is None:
        with tf.device('/device:CPU:0'):
          state[0] = tf.Variable(42)
      with tf.device('/device:CPU:0'):
        return x + state[0]

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(42, state[0].read_value())
    self.assertAllEqual(43, client.foo(1))
    state[0].assign(0)
    self.assertAllEqual(1, client.foo(1))
    server.shutdown()

  def test_wait_for_server(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)

    def create_client():
      result = ops.Client(address)
      return result

    with futures.ThreadPoolExecutor(max_workers=1) as executor:
      f = executor.submit(create_client)

      time.sleep(2)
      server.start()
      self.assertAllEqual(43, f.result().foo(42))
      server.shutdown()

  def test_wait_for_server2(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)

    def create_and_send():
      client = ops.Client(address)
      self.assertAllEqual(43, client.foo(42))

    with futures.ThreadPoolExecutor(max_workers=1) as executor:
      f = executor.submit(create_and_send)

      time.sleep(2)
      server.start()
      f.result()
      server.shutdown()

  def test_upvalue(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    a = tf.constant(2)

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x / a

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(21, client.foo(42))
    server.shutdown()

  def test_queue(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    q = tf.queue.FIFOQueue(1, [tf.int32], [()])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      if x.shape == (1,):
        q.enqueue_many([x])
      else:
        q.enqueue([x])
      return x

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    client.foo(42)
    self.assertAllEqual(42, q.dequeue())
    server.shutdown()

  def test_string(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.string)])
    def hello(x):
      return tf.strings.join([x, ' world'])

    server.bind(hello)
    server.start()

    client = ops.Client(address)
    self.assertAllEqual(b'hello world', client.hello('hello'))
    server.shutdown()

  def test_client_non_scalar_server_address(self):
    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                'server_address must be a scalar'):
      ops.Client(['localhost:8000', 'localhost:8001'])

  def test_server_non_vector_server_addresses(self):
    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                'server_addresses must be a vector'):
      ops.Server([['localhost:8000', 'localhost:8001']])

  def test_not_bound(self):
    address = self.get_unix_address()
    server = ops.Server([address])
    with self.assertRaisesRegex(tf.errors.UnavailableError,
                                'No function was bound'):
      server.start()

  def test_binding_function_twice(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[])
    def foo():
      return 42

    server.bind(foo)
    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                'Function \'foo\' was bound twice.'):
      server.bind(foo)

  def test_starting_twice(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)

    server.start()
    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                'Server is already started'):
      server.start()

  def test_invalid_number_of_arguments(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                'Expects 1 arguments, but 2 is provided'):
      client.foo([42, 43])
    server.shutdown()

  def test_invalid_type(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      return x + 1

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    with self.assertRaisesRegex(
        tf.errors.InvalidArgumentError,
        r'Expects arg\[0\] to be int32 but string is provided'):
      client.foo('foo')
    server.shutdown()

  def test_failing_function(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
    def foo(x):
      tf.assert_equal(1, x)  # Will fail.
      return x

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                'assertion failed'):
      client.foo(42)
    server.shutdown()

  def test_nests(self):
    address = self.get_unix_address()
    server = ops.Server([address])

    signature = (tf.TensorSpec([], tf.int32, name='arg1'),
                 Some(
                     tf.TensorSpec([], tf.int32, name='arg2'), [
                         tf.TensorSpec([], tf.int32, name='arg3'),
                         tf.TensorSpec([], tf.int32, name='arg4')
                     ]))

    @tf.function(input_signature=signature)
    def foo(*args):
      return tf.nest.map_structure(lambda t: t + 1, args)

    server.bind(foo)
    server.start()

    client = ops.Client(address)
    inputs = (1, Some(2, [3, 4]))
    expected_outputs = (2, Some(3, [4, 5]))
    outputs = client.foo(inputs)
    outputs = tf.nest.map_structure(lambda t: t.numpy(), outputs)
    tf.nest.assert_same_structure(expected_outputs, outputs)
    self.assertAllEqual(
        tf.nest.flatten(expected_outputs), tf.nest.flatten(outputs))
    server.shutdown()