# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# pylint: disable=invalid-name, unused-argument

import pytest

pytest.importorskip("ethosu.vela")
import numpy as np

import tvm
from tvm import relay
from tvm.relay.backend.contrib.ethosu import preprocess


def set_func_attr(func, compile_name, symbol_name):
    """
    Helper function to attach attributes to the external function.
    """
    func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
    func = func.with_attr("Inline", tvm.tir.IntImm("int32", 1))
    func = func.with_attr("Compiler", compile_name)
    func = func.with_attr("global_symbol", symbol_name)
    return func


def test_single_io():
    """
    This test will test the pass wont touch external functions that
    have a single input and a single output.
    """

    def create_graph():
        def create_external_func1(mod_, compiler_name, symbol_name):
            x_int = relay.var("x_int", shape=(10, 10))
            z0 = relay.nn.relu(x_int)
            f1 = relay.Function([x_int], z0)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()
        x = relay.var("x", shape=(10, 10))

        glb_symbol_f1, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        r = relay.Call(glb_symbol_f1, [x])
        main = relay.Function([x], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    mod = create_graph()
    exp = create_graph()
    mod = preprocess.preprocess_ext_io()(mod)
    assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)


def test_2ins_single_out():
    """
    The test is check two inputs and a single output of external function
    """

    def create_graph():
        def create_external_func1(mod_, compiler_name, symbol_name):
            x_int = relay.var("x_int", shape=(10, 10))
            w0_int = relay.var("w0_int", shape=(10, 10))
            z0 = relay.add(x_int, w0_int)

            f1 = relay.Function([x_int, w0_int], z0)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()

        x = relay.var("x", shape=(10, 10))
        w0 = relay.var("w0", shape=(10, 10))

        glb_symbol_f1, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        r = relay.Call(glb_symbol_f1, [x, w0])
        main = relay.Function([x, w0], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    def expected():
        def create_external_func1(mod_, compiler_name, symbol_name):
            ifms_int = relay.var("ifms_int", shape=[200])

            # splits
            (x_int_flat, w0_int_flat) = relay.split(ifms_int, [100])
            # reshapes
            x_int = relay.reshape(x_int_flat, newshape=(10, 10))
            w0_int = relay.reshape(w0_int_flat, newshape=(10, 10))

            z0 = relay.add(x_int, w0_int)
            f1 = relay.Function([ifms_int], z0)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()

        x = relay.var("x", shape=(10, 10))
        w0 = relay.var("w0", shape=(10, 10))

        # reshapes
        x_reshaped = relay.reshape(x, newshape=100)
        w0_reshaped = relay.reshape(w0, newshape=100)

        # concat
        ifms = relay.concatenate((x_reshaped, w0_reshaped), 0)

        glb_symbol_f1, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        r = relay.Call(glb_symbol_f1, [ifms])
        main = relay.Function([x, w0], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    mod = create_graph()
    exp = expected()
    mod = preprocess.preprocess_ext_io()(mod)
    assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)


def test_single_in_2outs():
    """
    The test is to check a single input and two outputs of external function
    """

    def create_graph():
        def create_external_func1(mod_, compiler_name, symbol_name):
            x_int = relay.var("x_int", shape=(10, 10))

            p0 = relay.nn.relu(x_int)
            q0 = relay.tanh(x_int)
            f1_o_tuple = relay.Tuple([p0, q0])

            f1 = relay.Function([x_int], f1_o_tuple)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()
        x = relay.var("x", shape=(10, 10))
        glb_symbol_f1, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        pq_tuple = relay.Call(glb_symbol_f1, [x])
        p0 = relay.TupleGetItem(pq_tuple, 0)
        q0 = relay.TupleGetItem(pq_tuple, 1)
        r = relay.concatenate((p0, q0), axis=0)
        main = relay.Function([x], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    def expected():
        def create_external_func1(mod_, compiler_name, symbol_name):
            x_int = relay.var("x_int", shape=(10, 10))

            p0 = relay.nn.relu(x_int)
            q0 = relay.tanh(x_int)

            # reshapes
            p0_reshaped = relay.reshape(p0, newshape=100)
            q0_reshaped = relay.reshape(q0, newshape=100)
            ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0)

            f1 = relay.Function([x_int], ofms)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()
        x = relay.var("x", shape=(10, 10))
        glb_symbol_f1, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        ofms = relay.Call(glb_symbol_f1, [x])

        # splits
        (p0_flat, q0_flat) = relay.split(ofms, [100])
        # reshapes
        p0_flat_reshaped = relay.reshape(p0_flat, newshape=(10, 10))
        q0_flat_reshaped = relay.reshape(q0_flat, newshape=(10, 10))
        # original output
        tuple_out = relay.Tuple([p0_flat_reshaped, q0_flat_reshaped])

        p0 = relay.TupleGetItem(tuple_out, 0)
        q0 = relay.TupleGetItem(tuple_out, 1)
        r = relay.concatenate((p0, q0), axis=0)
        main = relay.Function([x], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    mod = create_graph()
    exp = expected()
    mod = relay.transform.InferType()(mod)
    mod = preprocess.preprocess_ext_io()(mod)
    assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)


def test_4ins_2outs():
    """
    The test is to check a 4 inputs and two outputs of external function.
    This just stand as a general test for multiple ins/outs.
    """

    def create_graph():
        def create_external_func1(mod_, compiler_name, symbol_name):
            x_int = relay.var("x_int", shape=(10, 10))
            w0_int = relay.var("w0_int", shape=(10, 10))
            w1_int = relay.var("w1_int", shape=(10, 10))
            w2_int = relay.var("w2_int", shape=(10, 10))

            z0 = relay.add(x_int, w0_int)
            p0 = relay.subtract(z0, w1_int)
            q0 = relay.multiply(z0, w2_int)
            f1_o_tuple = relay.Tuple([p0, q0])

            f1 = relay.Function([x_int, w0_int, w1_int, w2_int], f1_o_tuple)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()

        x = relay.var("x", shape=(10, 10))
        w0 = relay.var("w0", shape=(10, 10))
        w1 = relay.var("w1", shape=(10, 10))
        w2 = relay.var("w2", shape=(10, 10))

        glb_symbol_f1, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        pq_tuple = relay.Call(glb_symbol_f1, [x, w0, w1, w2])

        p0 = relay.TupleGetItem(pq_tuple, 0)
        q0 = relay.TupleGetItem(pq_tuple, 1)
        r = relay.concatenate((p0, q0), axis=0)
        main = relay.Function([x, w0, w1, w2], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    def expected():
        def create_external_func1(mod_, compiler_name, symbol_name):
            ifms_int = relay.var("ifms_int", shape=[400])

            # splits
            (x_int_flat, w0_int_flat, w1_int_flat, w2_int_flat) = relay.split(
                ifms_int, [100, 200, 300]
            )
            # reshapes
            x_int = relay.reshape(x_int_flat, newshape=(10, 10))
            w0_int = relay.reshape(w0_int_flat, newshape=(10, 10))
            w1_int = relay.reshape(w1_int_flat, newshape=(10, 10))
            w2_int = relay.reshape(w2_int_flat, newshape=(10, 10))

            z0 = relay.add(x_int, w0_int)
            p0 = relay.subtract(z0, w1_int)
            q0 = relay.multiply(z0, w2_int)
            # f1_o_tuple = relay.Tuple([p0, q0])

            # reshapes
            p0_reshaped = relay.reshape(p0, newshape=100)
            q0_reshaped = relay.reshape(q0, newshape=100)
            ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0)

            f1 = relay.Function([ifms_int], ofms)
            f1 = set_func_attr(f1, compiler_name, symbol_name)
            glb_f1 = relay.GlobalVar(symbol_name)
            mod_[glb_f1] = f1
            mod_ = relay.transform.InferType()(mod_)
            return glb_f1, mod_

        mod = tvm.IRModule()

        x = relay.var("x", shape=(10, 10))
        w0 = relay.var("w0", shape=(10, 10))
        w1 = relay.var("w1", shape=(10, 10))
        w2 = relay.var("w2", shape=(10, 10))

        # reshapes
        x_reshaped = relay.reshape(x, newshape=100)
        w0_reshaped = relay.reshape(w0, newshape=100)
        w1_reshaped = relay.reshape(w1, newshape=100)
        w2_reshaped = relay.reshape(w2, newshape=100)

        # concat
        ifms = relay.concatenate((x_reshaped, w0_reshaped, w1_reshaped, w2_reshaped), 0)

        # call
        glb_func, mod = create_external_func1(mod, "ethos-u", "ethosu_0")
        ofms = relay.Call(glb_func, [ifms])

        # splits
        (p0_flat, q0_flat) = relay.split(ofms, [100])
        # reshapes
        p0_flat_reshaped = relay.reshape(p0_flat, newshape=(10, 10))
        q0_flat_reshaped = relay.reshape(q0_flat, newshape=(10, 10))
        # original output
        tuple_out = relay.Tuple([p0_flat_reshaped, q0_flat_reshaped])

        p0 = relay.TupleGetItem(tuple_out, 0)
        q0 = relay.TupleGetItem(tuple_out, 1)

        r = relay.concatenate((p0, q0), axis=0)
        main = relay.Function([x, w0, w1, w2], r)
        mod["main"] = main
        mod = relay.transform.InferType()(mod)
        return mod

    mod = create_graph()
    exp = expected()
    mod = preprocess.preprocess_ext_io()(mod)
    assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)


if __name__ == "__main__":
    pytest.main([__file__])