def test_set_idx_1d():
  idx = 10
  for vec in vecs:
    vec1, vec2 = vec.copy(), vec.copy()
    val = -vec[idx]
    vec2[idx] = val
    expect(set_idx_1d, [vec1, idx, val], vec2)

def test_constants_across_control_flow():
  testing_helpers.expect(const_across_control_flow, [True], 1)
  typed_fn = parakeet.typed_repr(const_across_control_flow, [True])
  assert len(typed_fn.body) == 1, "Fn body too long: " + str(typed_fn.body)
  stmt = typed_fn.body[0]
  assert isinstance(stmt, syntax.Return)
  assert isinstance(stmt.value, syntax.Const)

def test_always_true():
  testing_helpers.expect(always_true_branch, [], 0)
  typed_fn = parakeet.typed_repr(always_true_branch, [])
  assert len(typed_fn.body) == 1, "Fn body too long: " + str(typed_fn.body)
  stmt = typed_fn.body[0]
  assert isinstance(stmt, syntax.Return)
  assert isinstance(stmt.value, syntax.Const)

def test_set_idx_2d():
  i = 2
  j = 2
  for mat in matrices:
    mat1, mat2 = mat.copy(), mat.copy()
    val = -mat[i,j]
    mat2[i,j] = val
    expect(set_idx_2d, [mat1, i, j, val], mat2)

def test_set_idx_3d():
  i = 2
  j = 3
  k = 1
  for x in tensors:
    x1, x2 = x.copy(), x.copy()
    val = -x[i, j, k]
    x2[i, j, k] = val
    expect(set_idx_3d, [x1, i, j, k, val], x2)

def test_assign_slices():
  for m in matrices:
    m_expect = m.copy()
    m_input = m.copy()
    (i,j,k,l) = (0,2,0,4)
    (a,b,c,d) = (1,3,5,9)
    m_expect[i:j, k:l] = m_expect[a:b, c:d]
    expect(assign_slices, [m_input, (i,j,k,l), (a,b,c,d)], m_expect)
    expect(assign_slices, [m_input, (i,j,k,l), (a,b,c,d)], m_expect)

def test_copy_elimination():
  x = np.array([[1,2,3],[4,5,6]])
  expect(nested_add1, [x], x + 1.0)
  typed_fn = parakeet.typed_repr(nested_add1, [x])
  lowered = lowering.apply(typed_fn)
  n_loops = count_loops(lowered)
  n_expected = 3 if config.opt_loop_unrolling else 2
  assert n_loops <= n_expected, \
      "Too many loops generated! Expected at most 2, got %d" % n_loops

def all_tuples(f, unpack_args = True):
  """
  Given a function which should act as the identity, test it on multiple tuples
  """

  for t in [ints, mixed, nested2, nested2]:
    if unpack_args:
      expect(f, t, t)
    else:
      expect(f, [t], t)

def test_implicit_to_float():
  expect(implicit_to_float, [1], 1.5)
  expect(implicit_to_float, [True], 1.5)

def test_float_sum():
  testing_helpers.expect(my_sum, [float_vec], np.sum(float_vec))

def test_bool_sum():
  testing_helpers.expect(my_sum, [bool_vec], np.sum(bool_vec))

def test_if_true_const():
  expect(if_true_const, [], 1)

def test_int_sum():
  testing_helpers.expect(my_sum, [int_vec], np.sum(int_vec))

def test_assign_first_axis():
  for m in matrices:
    m_expect = m.copy()
    m_input = m.copy()
    m_expect[1] = m_expect[2]
    expect(assign_first_axis, [m_input, 1, 2], m_expect)

def test_varargs_return():
  expect(varargs_return, [1,2], (1,2))

def test_slice_first_axis_matrices():
  for m in matrices:
    expect(implicit_slice_first_axis, [m, 2], m[2])

def test_return_pair():
  expect(return_pair, [], (-1.0, 200))

def test_tuple_lhs_sum():
  tuples = [(True, 0, 1.0), (1, True, 0)]
  for t in tuples:
    expect(tuple_lhs_sum, [t], sum(t))

def test_implicit_to_bool():
  expect(implicit_to_bool, [1], 10)
  expect(implicit_to_bool, [2], 10)
  expect(implicit_to_bool, [0], -10)
  expect(implicit_to_bool, [1.0], 10)
  expect(implicit_to_bool, [2.0], 10)
  expect(implicit_to_bool, [0.0], -10)

def test_varargs_add():
  expect(varargs_add, [1,2], 3)