import tvm
from tvm import relay
from tvm.relay.ir_pass import (free_vars, free_type_vars,
                               bound_vars, bound_type_vars,
                               all_vars, all_type_vars)

def assert_vars_match(actual, expected):
    assert len(actual) == len(expected)
    for i in range(len(actual)):
        assert actual[i] == expected[i]


def test_free_vars():
    ty = relay.TensorType([], "int32")
    x = relay.Var("x", ty)
    fvx = free_vars(x)
    assert len(fvx) == 1
    assert fvx[0] == x
    v = relay.Constant(tvm.nd.array(10))

    let = relay.Let(x, v, x)
    fvx = free_vars(let)
    assert len(free_vars(let)) == 0
    f = relay.Function([x], x, ty)
    assert len(free_vars(f)) == 0


def test_free_vars_tuple():
    t = relay.Var('t')
    fv = free_vars(relay.Tuple([t, t]))
    assert len(fv) == 1
    assert fv[0] == t
    fv = free_vars(relay.TupleGetItem(t, 123))
    assert len(fv) == 1
    assert fv[0] == t


def test_free_type_vars():
    tp = relay.TypeVar("")
    ty = relay.TupleType([tp, relay.TensorType([], "int32")])
    x = relay.Var("x", ty)
    y = relay.Var("y")
    let = relay.Let(x, y, x)
    fvl = free_vars(let)
    assert len(fvl) == 1
    assert fvl[0] == y
    ftvl = free_type_vars(let)
    assert len(ftvl) == 1
    assert ftvl[0] == tp


def test_bound_vars():
    x = relay.Var("x")
    y = relay.Var("y")
    z = relay.Var("z")
    a = relay.Var("a")

    f1 = relay.Function([x, y, z], relay.Let(a, x, relay.Tuple([])))
    assert_vars_match(bound_vars(f1), [x, y, z, a])

    tup = relay.Tuple([x, y, z, a])
    assert len(bound_vars(tup)) == 0

    f2 = relay.Function([x, y], relay.Tuple([x, y, z, a]))
    assert_vars_match(bound_vars(f2), [x, y])


def test_bound_type_vars():
    a = relay.TypeVar("a")
    b = relay.TypeVar("b")
    c = relay.TypeVar("c")

    ft1 = relay.FuncType([a], b, [a, b])
    bound_ft1 = bound_type_vars(ft1)
    assert_vars_match(bound_type_vars(ft1), [a, b])

    ft2 = relay.FuncType([], c, [a])
    assert_vars_match(bound_type_vars(ft2), [a])

    tup_ty = relay.TupleType([a, b, c])
    assert len(bound_type_vars(tup_ty)) == 0

    f1 = relay.Function([], relay.Tuple([]), type_params=[a, b])
    assert_vars_match(bound_type_vars(f1), [a, b])

    f2 = relay.Function([], relay.Tuple([]), c)
    assert len(bound_type_vars(f2)) == 0

    x = relay.Var("x", a)
    let1 = relay.Let(x, relay.Tuple([]), x)
    assert len(bound_type_vars(let1)) == 0

    let2 = relay.Let(x, relay.Function([], relay.Tuple([]), type_params=[b, c]), x)
    assert_vars_match(bound_type_vars(let2), [b, c])


def test_all_vars():
    x = relay.Var("x")
    y = relay.Var("y")
    z = relay.Var("z")

    f1 = relay.Function([x, y], z)
    assert_vars_match(all_vars(f1), [x, y, z])

    f2 = relay.Function([x], relay.Let(y, relay.Tuple([]), z))
    assert_vars_match(all_vars(f2), [x, y, z])

    f3 = relay.Function([x], relay.Tuple([y, z]))
    assert_vars_match(all_vars(f3), [x, y, z])

    tup = relay.Tuple([x, y, z])
    assert_vars_match(all_vars(tup), [x, y, z])


def test_all_type_vars():
    a = relay.TypeVar("a")
    b = relay.TypeVar("b")
    c = relay.TypeVar("c")

    ft1 = relay.FuncType([b], c, [a])
    assert_vars_match(all_type_vars(ft1), [a, b, c])

    ft2 = relay.FuncType([], relay.TupleType([a, b, c]), [])
    assert_vars_match(all_type_vars(ft2), [a, b, c])

    w = relay.Var("w")
    x = relay.Var("x", a)
    y = relay.Var("y", b)
    z = relay.Var("z", c)
    
    f1 = relay.Function([x], y, b, [a])
    assert_vars_match(all_type_vars(f1), [a, b])

    f2 = relay.Function([x], relay.Let(y, x, z))
    assert_vars_match(all_type_vars(f2), [a, b, c])

    f3 = relay.Function([], relay.Tuple([x, y, z]), ret_type=relay.TupleType([a, b, c]))
    assert_vars_match(all_type_vars(f3), [a, b, c])

    f4 = relay.Function([w], relay.Tuple([]), type_params=[a, b, c])
    assert_vars_match(all_type_vars(f4), [a, b, c])

    f5 = relay.Function([w], w)
    assert len(all_type_vars(f5)) == 0
