# 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.
import numpy as np
import pytest
from unittest import mock

import tvm
import tvm.testing
from tvm import relay
from tvm.relay import transform
from tvm.contrib.target import coreml as _coreml

requires_coremltools = tvm.testing.requires_package("coremltools")


def _has_xcode():
    try:
        tvm.contrib.xcode.xcrun([])
        return True
    except FileNotFoundError:
        pass

    return False


def _create_graph():
    shape = (10, 10)
    mod = tvm.IRModule()

    x = relay.var("x", shape=shape)
    y = relay.var("y", shape=shape)
    z = x + x
    p = y * y
    func = relay.Function([x, y], p - z)
    mod["main"] = func

    return mod


def _create_graph_annotated():
    shape = (10, 10)
    target = "coremlcompiler"
    mod = tvm.IRModule()

    # function 0
    f0_i0 = relay.var(target + "_0_i0", shape=shape)
    func0 = relay.Function([f0_i0], f0_i0 * f0_i0)

    func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
    func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
    func0 = func0.with_attr("Compiler", target)
    func0 = func0.with_attr("global_symbol", target + "_0")
    gv0 = relay.GlobalVar(target + "_0")
    mod[gv0] = func0

    # function 2
    f2_i0 = relay.var(target + "_2_i0", shape=shape)
    func2 = relay.Function([f2_i0], f2_i0 + f2_i0)

    func2 = func2.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
    func2 = func2.with_attr("Inline", tvm.tir.IntImm("int32", 1))
    func2 = func2.with_attr("Compiler", target)
    func2 = func2.with_attr("global_symbol", target + "_2")
    gv2 = relay.GlobalVar(target + "_2")
    mod[gv2] = func2
    mod = relay.transform.InferType()(mod)

    # body
    x = relay.var("x", shape=shape)
    y = relay.var("y", shape=shape)
    func = relay.Function([x, y], gv0(y) - gv2(x))
    mod["main"] = func
    mod = relay.transform.InferType()(mod)

    return mod


@pytest.mark.xfail(
    reason="Currently failing test.  See tracking issue https://github.com/apache/tvm/issues/8901"
)
@tvm.testing.uses_gpu
@requires_coremltools
def test_annotate():
    mod = _create_graph()
    mod = transform.AnnotateTarget("coremlcompiler")(mod)
    mod = transform.PartitionGraph()(mod)

    expected = _create_graph_annotated()
    assert tvm.ir.structural_equal(mod, expected, map_free_vars=True)


@pytest.mark.skipif(not _has_xcode(), reason="Xcode is not available")
@tvm.testing.uses_gpu
@requires_coremltools
def test_compile_and_run():
    dev = tvm.cpu()
    target = "llvm"
    tol = 1e-3

    with relay.build_config(opt_level=3):
        lib = relay.build(_create_graph_annotated(), target=target)
    m = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))

    shape = (10, 10)
    x_data = np.random.rand(*shape).astype("float32")
    y_data = np.random.rand(*shape).astype("float32")

    m.set_input("x", x_data)
    m.set_input("y", y_data)
    m.run()
    out = tvm.nd.empty(shape, device=dev)
    out = m.get_output(0, out)

    expected = (y_data * y_data) - (x_data + x_data)
    tvm.testing.assert_allclose(out.numpy(), expected, rtol=tol, atol=tol)


@mock.patch("tvm.contrib.coreml_runtime.create")
@mock.patch("tvm.contrib.xcode.compile_coreml")
def _construct_model(func, m1, m2):
    mod = tvm.IRModule()
    mod["main"] = func
    mod = transform.AnnotateTarget("coremlcompiler")(mod)
    mod = transform.PartitionGraph()(mod)

    fcompile = tvm._ffi.get_global_func("relay.ext.coremlcompiler")

    for var, func in mod.functions.items():
        if func.attrs and "Compiler" in func.attrs and func.attrs["Compiler"] == "coremlcompiler":
            fcompile(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_add():
    shape = (10, 10)
    x = relay.var("x", shape=shape)
    y = x + x
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_multiply():
    shape = (10, 10)
    x = relay.var("x", shape=shape)
    y = x * x
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_clip():
    shape = (10, 10)
    x = relay.var("x", shape=shape)
    y = relay.clip(x, a_min=0.0, a_max=1.0)
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_batch_flatten():
    shape = (10, 10, 10)
    x = relay.var("x", shape=shape)
    y = relay.nn.batch_flatten(x)
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_expand_dims():
    shape = (10, 10)
    x = relay.var("x", shape=shape)
    y = relay.expand_dims(x, axis=0)
    func = relay.Function([x], y)
    _construct_model(func)

    y = relay.expand_dims(x, axis=-1)
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_relu():
    shape = (10, 10)
    x = relay.var("x", shape=shape)
    y = relay.nn.relu(x)
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_softmax():
    shape = (10, 10)
    x = relay.var("x", shape=shape)
    y = relay.nn.softmax(x, axis=1)
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_conv2d():
    x = relay.var("x", shape=(1, 3, 224, 224))
    w = relay.const(np.zeros((16, 3, 3, 3), dtype="float32"))
    y = relay.nn.conv2d(x, w, strides=[2, 2], padding=[1, 1, 1, 1], kernel_size=[3, 3])
    func = relay.Function([x], y)
    _construct_model(func)


@tvm.testing.uses_gpu
@requires_coremltools
def test_global_avg_pool2d():
    shape = (10, 10, 10, 10)
    x = relay.var("x", shape=shape)
    y = relay.nn.global_avg_pool2d(x)
    func = relay.Function([x], y)
    _construct_model(func)


if __name__ == "__main__":
    test_annotate()
    test_compile_and_run()
    test_add()
    test_multiply()
    test_clip()
    test_expand_dims()
    test_relu()
    test_batch_flatten()
    test_softmax()
    test_conv2d()
    test_global_avg_pool2d()