import ray

import numpy as np

# Test simple functionality

@ray.remote(num_return_vals=2)
def handle_int(a, b):
  return a + 1, b + 1

# Test aliasing

@ray.remote()
def test_alias_f():
  return np.ones([3, 4, 5])

@ray.remote()
def test_alias_g():
  return test_alias_f.remote()

@ray.remote()
def test_alias_h():
  return test_alias_g.remote()

# Test timing

@ray.remote()
def empty_function():
  pass

@ray.remote()
def trivial_function():
  return 1

# Test keyword arguments

@ray.remote()
def keyword_fct1(a, b="hello"):
  return "{} {}".format(a, b)

@ray.remote()
def keyword_fct2(a="hello", b="world"):
  return "{} {}".format(a, b)

@ray.remote()
def keyword_fct3(a, b, c="hello", d="world"):
  return "{} {} {} {}".format(a, b, c, d)

# Test variable numbers of arguments

@ray.remote()
def varargs_fct1(*a):
  return " ".join(map(str, a))

@ray.remote()
def varargs_fct2(a, *b):
  return " ".join(map(str, b))

try:
  @ray.remote()
  def kwargs_throw_exception(**c):
    return ()
  kwargs_exception_thrown = False
except:
  kwargs_exception_thrown = True

try:
  @ray.remote()
  def varargs_and_kwargs_throw_exception(a, b="hi", *c):
    return "{} {} {}".format(a, b, c)
  varargs_and_kwargs_exception_thrown = False
except:
  varargs_and_kwargs_exception_thrown = True

# test throwing an exception

@ray.remote()
def throw_exception_fct1():
  raise Exception("Test function 1 intentionally failed.")

@ray.remote()
def throw_exception_fct2():
  raise Exception("Test function 2 intentionally failed.")

@ray.remote(num_return_vals=3)
def throw_exception_fct3(x):
  raise Exception("Test function 3 intentionally failed.")

# test Python mode

@ray.remote()
def python_mode_f():
  return np.array([0, 0])

@ray.remote()
def python_mode_g(x):
  x[0] = 1
  return x

# test no return values

@ray.remote()
def no_op():
  pass

@ray.remote()
def no_op_fail():
  return 0

# test wrong return types

@ray.remote()
def test_return1():
  return 0.0

@ray.remote(num_return_vals=2)
def test_return2():
  return 2.0, 3.0

class TestClass(object):
  def __init__(self):
    self.a = 5

@ray.remote()
def test_unknown_type():
  return TestClass()