import nnvm
import nnvm.symbol as sym
import nnvm.graph as graph
from nnvm.compiler import graph_attr

def correct_layout(g, layout=None):
    if isinstance(g, nnvm.symbol.Symbol):
        g = graph.create(g)
    if layout:
        graph_attr.set_layout_inputs(g, layout)
    g = g.apply("CorrectLayout")
    ldict = {}
    vlayout = g.json_attr("layout")
    entry_ptr = g.index.entry_ptr
    for i, n in enumerate(g.index.nodes):
        begin, end = entry_ptr[i], entry_ptr[i + 1]
        ldict[n["name"]] = vlayout[begin:end]
    return g, ldict


# Level 1
def test_dense():
    x = sym.Variable("data", shape=(10, 20))
    y = sym.dense(x, units=30, name="fc")
    g, ldict = correct_layout(y, "HW")
    assert(ldict["data"][0] == "HW")
    assert(ldict["fc"][0] == "HW")
    assert(ldict["fc_bias"][0] == "__undef__")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, "HW16w")
    assert(ldict["data"][0] == "HW16w")
    assert(ldict["data_HW"][0] == "HW")
    assert(ldict["fc"][0] == "HW")
    assert(ldict["fc_bias"][0] == "__undef__")


def test_matmul():
    a = sym.Variable("a", shape=(10, 20))
    b = sym.Variable("b", shape=(20, 30))
    c = sym.matmul(a, b, name="matmul")
    g, ldict = correct_layout(c, {"a" : "HW", "b" : "WC"})
    assert(ldict["a"][0] == "HW")
    assert(ldict["b"][0] == "WC")
    assert(ldict["matmul"][0] == "HC")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, {"a" : "HW16w", "b" : "WC16c"})
    assert(ldict["a"][0] == "HW16w")
    assert(ldict["a_HW"][0] == "HW")
    assert(ldict["b"][0] == "WC16c")
    assert(ldict["b_WC"][0] == "WC")
    assert(ldict["matmul"][0] == "HC")
    a = sym.Variable("a", shape=(20, 10))
    c = sym.matmul(a, b, name="matmul", transpose_a=True)
    g, ldict = correct_layout(c, {"a" : "HW", "b" : "HC"})
    assert(ldict["a"][0] == "HW")
    assert(ldict["b"][0] == "HC")
    assert(ldict["matmul"][0] == "WC")
    b = sym.Variable("b", shape=(30, 20))
    c = sym.matmul(a, b, name="matmul", transpose_b=True)
    g, ldict = correct_layout(c, {"a" : "HW", "b" : "CW"})
    assert(ldict["a"][0] == "HW")
    assert(ldict["b"][0] == "CW")
    assert(ldict["matmul"][0] == "HC")
    a = sym.Variable("a", shape=(20, 10))
    b = sym.Variable("b", shape=(30, 20))
    c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
    g, ldict = correct_layout(c, {"a" : "HW", "b" : "CH"})
    assert(ldict["a"][0] == "HW")
    assert(ldict["b"][0] == "CH")
    assert(ldict["matmul"][0] == "WC")


def test_concatenate():
    x1 = sym.Variable("x", shape=(10, 20))
    x2 = sym.Variable("y", shape=(10, 30))
    z = sym.concatenate(x1, x2, name="concat")
    g, ldict = correct_layout(z, {"x": "HW", "y": "HW"})
    assert(ldict["x"][0] == "HW")
    assert(ldict["y"][0] == "HW")
    assert(ldict["concat"][0] == "HW")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, {"x": "HW16w", "y": "HW16w"})
    assert(ldict["x"][0] == "HW16w")
    assert(ldict["y"][0] == "HW16w")
    assert(ldict["concat"][0] == "HW16w")

    x1 = sym.Variable("x", shape=(10, 20, 60))
    x2 = sym.Variable("y", shape=(10, 20, 40))
    z = sym.concatenate(x1, x2, axis=2, name="concat")
    g, ldict = correct_layout(z, {"x": "H20wW", "y": "H20wW"})
    assert(ldict["x"][0] == "H20wW")
    assert(ldict["y"][0] == "H20wW")
    assert(ldict["concat"][0] == "H20wW")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, {"x": "HW", "y": "HW"})
    assert(ldict["x_H20wW"][0] == "H20wW")
    assert(ldict["x_H20wW"][0] == "H20wW")
    assert(ldict["concat"][0] == "H20wW")


def test_expand_dims():
    x = sym.Variable("x", shape=(10, 20))
    y = sym.expand_dims(x, axis=1, name="y")
    g, ldict = correct_layout(y, "HW")
    assert(ldict["x"][0] == "HW")
    assert(ldict["y"][0] == "__undef__")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, "HW16w")
    assert(ldict["x"][0] == "HW16w")
    assert(ldict["x_HW"][0] == "HW")
    assert(ldict["y"][0] == "__undef__")


def test_split():
    x = sym.Variable("x", shape=(10, 20))
    y = sym.split(x, indices_or_sections=[11], name="y")
    g, ldict = correct_layout(y, "HW")
    assert(ldict["x"][0] == "HW")
    assert(ldict["y"][0] == "__undef__")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, "HW16w")
    assert(ldict["x"][0] == "HW16w")
    assert(ldict["x_HW"][0] == "HW")
    assert(ldict["y"][0] == "__undef__")


def test_batchnorm():
    x = sym.Variable("data", shape=(10, 20, 30, 40))
    y = sym.batch_norm(x, axis=1, epsilon=2e-5, name="bn")
    g, ldict = correct_layout(y, "NCHW")
    assert(ldict["data"][0] == "NCHW")
    assert(ldict["bn"][0] == "NCHW")
    assert(ldict["bn"][1] == "C")
    assert(ldict["bn"][2] == "C")
    assert(ldict["bn_beta"][0] == "C")
    assert(ldict["bn_gamma"][0] == "C")
    assert(ldict["bn_moving_mean"][0] == "C")
    assert(ldict["bn_moving_var"][0] == "C")
    # batch_norm can deal with sub-dim of C at the last dim.
    g, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["data"][0] == "NCHW16c")
    assert(ldict["bn"][0] == "NCHW16c")
    assert(ldict["bn"][1] == "C16c")
    assert(ldict["bn"][2] == "C16c")
    assert(ldict["bn_beta"][0] == "C")
    assert(ldict["bn_beta_C16c"][0] == "C16c")
    assert(ldict["bn_gamma"][0] == "C")
    assert(ldict["bn_gamma_C16c"][0] == "C16c")
    assert(ldict["bn_moving_mean"][0] == "C")
    assert(ldict["bn_moving_mean_C16c"][0] == "C16c")
    assert(ldict["bn_moving_var"][0] == "C")
    assert(ldict["bn_moving_var_C16c"][0] == "C16c")
    # but for other layout, it does a layout transform for data
    g, ldict = correct_layout(g, "NCH16cW")
    assert(ldict["data"][0] == "NCH16cW")
    assert(ldict["data_NCHW16c"][0] == "NCHW16c")
    assert(ldict["bn"][0] == "NCHW16c")
    assert(ldict["bn"][1] == "C16c")
    assert(ldict["bn"][2] == "C16c")
    assert(ldict["bn_beta"][0] == "C")
    assert(ldict["bn_beta_C16c"][0] == "C16c")
    assert(ldict["bn_gamma"][0] == "C")
    assert(ldict["bn_gamma_C16c"][0] == "C16c")
    assert(ldict["bn_moving_mean"][0] == "C")
    assert(ldict["bn_moving_mean_C16c"][0] == "C16c")
    assert(ldict["bn_moving_var"][0] == "C")
    assert(ldict["bn_moving_var_C16c"][0] == "C16c")


def test_flatten():
    x = sym.Variable("x", shape=(10, 20, 10, 10))
    y = sym.flatten(x, name="y")
    g, ldict = correct_layout(y, "NCHW")
    assert(ldict["x"][0] == "NCHW")
    assert(ldict["y"][0] == "__undef__")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["x"][0] == "NCHW16c")
    assert(ldict["x_NCHW"][0] == "NCHW")
    assert(ldict["y"][0] == "__undef__")


def test_softmax():
    x = sym.Variable("x", shape=(10, 20, 10, 10))
    y = sym.softmax(x, name="y")
    g, ldict = correct_layout(y, "NCHW")
    assert(ldict["x"][0] == "NCHW")
    assert(ldict["y"][0] == "NCHW")
    # second pass will insert layout transform
    _, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["x"][0] == "NCHW16c")
    assert(ldict["x_NCHW"][0] == "NCHW")
    assert(ldict["y"][0] == "NCHW")


# Level 2
def test_conv2d():
    x = sym.Variable("data", shape=(1, 32, 512, 512))
    y = sym.conv2d(x, name="conv", channels=12,
                   kernel_size=(3,3), padding=(1,1), layout="NCHW")
    _, ldict = correct_layout(y)
    assert(ldict["data"][0] == "NCHW")
    assert(ldict["conv_weight"][0] == "OIHW")
    assert(ldict["conv_bias"][0] == "C")
    assert(ldict["conv"][0] == "NCHW")
    y = sym.conv2d(x, name="conv", channels=12,
                   kernel_size=(3,3), padding=(1,1), layout="NCHW16c",
                   kernel_layout="OIHW16i16o", out_layout="NCHW8c")
    _, ldict = correct_layout(y)
    assert(ldict["data"][0] == "NCHW16c")
    assert(ldict["conv_weight"][0] == "OIHW16i16o")
    assert(ldict["conv_bias"][0] == "C8c")
    assert(ldict["conv"][0] == "NCHW8c")
    y = sym.conv2d(x, name="conv", channels=12,
                   kernel_size=(3,3), padding=(1,1), layout="N16cHWC")
    _, ldict = correct_layout(y)
    assert(ldict["data"][0] == "N16cHWC")
    assert(ldict["conv_weight"][0] == "OIHW")
    assert(ldict["conv_bias"][0] == "16cC")
    assert(ldict["conv"][0] == "N16cHWC")


def test_conv2d_transpose():
    x = sym.Variable("data", shape=(1, 32, 512, 512))
    y = sym.conv2d_transpose(x, name="conv", channels=12,
                             kernel_size=(3,3), padding=(1,1), layout="NCHW")
    _, ldict = correct_layout(y)
    assert(ldict["data"][0] == "NCHW")
    assert(ldict["conv_weight"][0] == "OIHW")
    assert(ldict["conv_bias"][0] == "C")
    assert(ldict["conv"][0] == "NCHW")


def test_max_pool2d():
    x = sym.Variable("data", shape=(1, 32, 512, 512))
    y = sym.max_pool2d(x, name="pool", pool_size=(3,3),
                       padding=(1,1), layout="NCHW")
    g, ldict = correct_layout(y)
    assert(ldict["data"][0] == "NCHW")
    assert(ldict["pool"][0] == "NCHW")
    # if index of H and W remain the same,
    # pool2d does not convert the layout.
    g, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["data"][0] == "NCHW16c")
    assert(ldict["pool"][0] == "NCHW16c")
    # for other layout it requires a layout transform.
    g, ldict = correct_layout(g, "NHWC")
    assert(ldict["data"][0] == "NHWC")
    assert(ldict["data_NCHW"][0] == "NCHW")
    assert(ldict["pool"][0] == "NCHW")


def test_global_pool2d():
    x = sym.Variable("data", shape=(1, 32, 512, 512))
    y = sym.global_max_pool2d(x, name="pool", layout="NCHW")
    g, ldict = correct_layout(y)
    assert(ldict["data"][0] == "NCHW")
    assert(ldict["pool"][0] == "NCHW")
    # if index of H and W remain the same,
    # pool2d does not convert the layout.
    g, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["data"][0] == "NCHW16c")
    assert(ldict["pool"][0] == "NCHW16c")
    # for other layout it requires a layout transform.
    g, ldict = correct_layout(g, "NHWC")
    assert(ldict["data"][0] == "NHWC")
    assert(ldict["data_NCHW"][0] == "NCHW")
    assert(ldict["pool"][0] == "NCHW")


# Level 3
def test_reshape():
    x = sym.Variable("x", shape=(4,))
    y = sym.reshape(x, shape=(2,2), name="y")
    g, ldict = correct_layout(y, "C")
    assert(ldict["x"][0] == "C")
    assert(ldict["y"][0] == "__undef__")
    # second pass will insert layout transform
    g, ldict = correct_layout(g, "C16c")
    assert(ldict["x"][0] == "C16c")
    assert(ldict["x_C"][0] == "C")
    assert(ldict["y"][0] == "__undef__")


def test_transpose():
    x = sym.Variable("x", shape=(1, 32, 512, 512))
    y = sym.transpose(x, name="y", axes=(0, 2, 3, 1))
    g, ldict = correct_layout(y, "NCHW")
    assert(ldict["x"][0] == "NCHW")
    assert(ldict["y"][0] == "NHWC")
    # second pass will insert layout transform
    g, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["x"][0] == "NCHW16c")
    assert(ldict["x_NCHW"][0] == "NCHW")
    assert(ldict["y"][0] == "NHWC")


def test_broadcast_to():
    x = sym.Variable("x", shape=(4, 1))
    y = sym.broadcast_to(x, shape=(0, 4), name="y")
    g, ldict = correct_layout(y, "HW")
    assert(ldict["x"][0] == "HW")
    assert(ldict["y"][0] == "__undef__")
    # second pass will insert layout transform
    g, ldict = correct_layout(g, "HW16h")
    assert(ldict["x"][0] == "HW16h")
    assert(ldict["x_HW"][0] == "HW")
    assert(ldict["y"][0] == "__undef__")


def test_broadcast_binary():
    x = sym.Variable("x", shape=(1, 16, 512, 512))
    y = sym.Variable("y", shape=(16, 512, 512))
    z = sym.broadcast_add(x, y, name="z")
    g, ldict = correct_layout(z, {"x": "NCHW", "y": "CHW"})
    assert(ldict["x"][0] == "NCHW")
    assert(ldict["y"][0] == "CHW")
    assert(ldict["z"][0] == "NCHW")
    # prior to keep the left layout if they do not match.
    g, ldict = correct_layout(g, {"x": "NCHW16c", "y": "CHW"})
    assert(ldict["x"][0] == "NCHW16c")
    assert(ldict["y"][0] == "CHW")
    assert(ldict["y_CHW16c"][0] == "CHW16c")
    assert(ldict["z"][0] == "NCHW16c")
    # broadcast_add(HCW16c, N16nCH16cW)
    g, ldict = correct_layout(z, {"x": "HCW16c", "y": "N16nCH16cW"})
    assert(ldict["x"][0] == "HCW16c")
    assert(ldict["y"][0] == "N16nCH16cW")
    assert(ldict["x_CH16cW"][0] == "CH16cW")
    assert(ldict["z"][0] == "N16nCH16cW")


def test_reduce():
    x = sym.Variable("x", shape=(1, 16, 512, 512))
    y = sym.sum(x, name="y", axis=1)
    g, ldict = correct_layout(y, "NCHW")
    assert(ldict["x"][0] == "NCHW")
    assert(ldict["y"][0] == "__undef__")
    # second pass will insert layout transform
    g, ldict = correct_layout(g, "NCHW16c")
    assert(ldict["x"][0] == "NCHW16c")
    assert(ldict["x_NCHW"][0] == "NCHW")
    assert(ldict["y"][0] == "__undef__")


if __name__ == "__main__":
    test_dense()
    test_matmul()
    test_concatenate()
    test_expand_dims()
    test_split()
    test_batchnorm()
    test_flatten()
    test_softmax()
    test_conv2d()
    test_conv2d_transpose()
    test_max_pool2d()
    test_global_pool2d()
    test_reshape()
    test_transpose()
    test_broadcast_to()
    test_broadcast_binary()
    test_reduce()
