# 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.

# This test checks if dynamic loading of library into MXNet is successful

import os
import platform
import mxnet as mx
import numpy as np
from mxnet import nd
from mxnet.gluon import nn
from mxnet.base import MXNetError
from mxnet.test_utils import download, is_cd_run, assert_almost_equal, default_device
import pytest

base_path = os.path.join(os.path.dirname(__file__), "../../..")
def check_platform(supported_platforms=['x86_64', 'AMD64']):
    return platform.machine() not in supported_platforms

@pytest.mark.skipif(check_platform(), reason="not all machine types supported")
@pytest.mark.skipif(is_cd_run(), reason="continuous delivery run - ignoring test")
def test_custom_op():
    # possible places to find library file
    if (os.name=='posix'):
        lib = 'libcustomop_lib.so'
        if os.path.exists(lib):
            fname = lib
        elif os.path.exists(os.path.join(base_path,'build/'+lib)):
            fname = os.path.join(base_path,'build/'+lib)
        else:
            raise MXNetError(f"library {lib} not found ")
    elif (os.name=='nt'):
        lib = 'libcustomop_lib.dll'
        if os.path.exists('windows_package\\lib\\'+lib):
            fname = 'windows_package\\lib\\'+lib
        else:
            raise MXNetError(f"library {lib} not found ")

    fname = os.path.abspath(fname)
    # load the library containing gemm custom operators
    mx.library.load(fname)

    # test symbol 2D gemm custom operators
    s = mx.sym.Variable('s')
    t = mx.sym.Variable('t')
    c = mx.sym.my_gemm(s,t)
    d = mx.sym.state_gemm(s,t)
    # baseline gemm from MXNet
    base = mx.sym.linalg.gemm2(s,t)

    # get some random input matrices
    dim_n, dim_k, dim_m = tuple(np.random.randint(1, 5, size=3))
    mat1 = mx.nd.random.uniform(-10, 10, shape=(dim_n, dim_k), ctx=mx.cpu())
    mat2 = mx.nd.random.uniform(-10, 10, shape=(dim_k, dim_m), ctx=mx.cpu())

    # intermediate ndarrays to be populated by gradient compute
    in_grad1 = [mx.nd.empty((dim_n,dim_k),ctx=mx.cpu()),mx.nd.empty((dim_k,dim_m),ctx=mx.cpu())]
    in_grad2 = [mx.nd.empty((dim_n,dim_k),ctx=mx.cpu()),mx.nd.empty((dim_k,dim_m),ctx=mx.cpu())]
    in_grad_base = [mx.nd.empty((dim_n,dim_k),ctx=mx.cpu()),mx.nd.empty((dim_k,dim_m),ctx=mx.cpu())]

    exe1 = c._bind(ctx=mx.cpu(),args={'s':mat1,'t':mat2},args_grad=in_grad1)
    exe2 = d._bind(ctx=mx.cpu(),args={'s':mat1,'t':mat2},args_grad=in_grad2)
    exe_base = base._bind(ctx=mx.cpu(),args={'s':mat1,'t':mat2},args_grad=in_grad_base)

    out1 = exe1.forward()
    out2 = exe2.forward()
    # test stateful operator by calling it multiple times
    out2 = exe2.forward()
    out_base = exe_base.forward()

    # check that forward compute matches one executed by MXNet
    assert_almost_equal(out_base[0].asnumpy(), out1[0].asnumpy(), rtol=1e-3, atol=1e-3)
    assert_almost_equal(out_base[0].asnumpy(), out2[0].asnumpy(), rtol=1e-3, atol=1e-3)

    # random output grad ndarray for gradient update
    out_grad = mx.nd.ones((dim_n, dim_m), ctx=mx.cpu())
    exe1.backward([out_grad])
    exe2.backward([out_grad])
    exe_base.backward([out_grad])

    # check that gradient compute matches one executed by MXNet
    assert_almost_equal(in_grad_base[0].asnumpy(), in_grad1[0].asnumpy(), rtol=1e-3, atol=1e-3)
    assert_almost_equal(in_grad_base[0].asnumpy(), in_grad2[0].asnumpy(), rtol=1e-3, atol=1e-3)

@pytest.mark.skipif(check_platform(), reason="not all machine types supported")
@pytest.mark.skipif(is_cd_run(), reason="continuous delivery run - ignoring test")
def test_subgraph():
    # possible places to find library file
    if (os.name=='posix'):
        lib = 'libsubgraph_lib.so'
        if os.path.exists(lib):
            # plain make build, when run in the CI
            fname = lib
        elif os.path.exists(os.path.join(base_path, 'build/'+lib)):
            # plain cmake build when run in the CI
            fname = os.path.join(base_path, 'build/'+lib)
        else:
            raise MXNetError(f"library {lib} not found ")
    elif (os.name=='nt'):
        lib = 'libsubgraph_lib.dll'
        if os.path.exists('windows_package\\lib\\'+lib):
            # plain make build, when run in the CI
            fname = 'windows_package\\lib\\'+lib
        else:
            # plain cmake build when run in the CI
            raise MXNetError(f"library {lib} not found ")

    fname = os.path.abspath(fname)
    mx.library.load(fname)

    # test simple graph with add, exp and log operators, library supports exp/log
    a = mx.sym.var('a')
    b = mx.sym.var('b')
    c = a + b
    d = mx.sym.exp(c)
    sym = mx.sym.log(d)

    args = {'a':mx.nd.ones((3,2),ctx=mx.cpu()), 'b':mx.nd.ones((3,2),ctx=mx.cpu())}

    # baseline - regular execution in MXNet
    exe = sym._bind(ctx=mx.cpu(), args=args)
    out = exe.forward()

    # without propogating shapes/types, passing a custom option to subgraph prop "myOpt"
    # should not create subgraph since subgraph prop requires type info
    mysym1 = sym.optimize_for("myProp", myOpt='yello')
    exe1 = mysym1._bind(ctx=mx.cpu(), args=args)
    out1 = exe1.forward()
    # check that result matches one executed by MXNet
    assert_almost_equal(out[0].asnumpy(), out1[0].asnumpy(), rtol=1e-3, atol=1e-3)

    # with propogating shapes/types, rejecting subgraph
    # this tests creating the subgraph and having the subgraph prop reject it
    mysym2 = sym.optimize_for("myProp", args, reject=True)
    exe2 = mysym2._bind(ctx=mx.cpu(), args=args)
    out2 = exe2.forward()
    # check that result matches one executed by MXNet
    assert_almost_equal(out[0].asnumpy(), out2[0].asnumpy(), rtol=1e-3, atol=1e-3)

    # with propogating shapes/types
    mysym3 = sym.optimize_for("myProp",args)
    exe3 = mysym3._bind(ctx=mx.cpu(), args=args)
    out3 = exe3.forward()
    # check that result matches one executed by MXNet
    assert_almost_equal(out[0].asnumpy(), out3[0].asnumpy(), rtol=1e-3, atol=1e-3)

    # Gluon Hybridize partitioning with shapes/types
    sym_block = nn.SymbolBlock(sym, [a,b])
    sym_block.initialize()
    sym_block.optimize_for(mx.nd.ones((3,2)),mx.nd.ones((3,2)),backend='myProp')
    out4 = sym_block(mx.nd.ones((3,2)),mx.nd.ones((3,2)))
    # check that result matches one executed by MXNet
    assert_almost_equal(out[0].asnumpy(), out4[0].asnumpy(), rtol=1e-3, atol=1e-3)

    # Gluon Hybridize partitioning with sym.var
    sym_block2 = nn.SymbolBlock(sym, [a,b])
    sym_block2.initialize()
    a_var = mx.sym.var('a',shape=(3,2))
    b_var = mx.sym.var('b',shape=(3,2))
    sym_block2.optimize_for(a_var, b_var, backend='myProp')

    # Gluon Hybridize partitioning with shapes/types
    sym_block3 = nn.SymbolBlock(sym, [a,b])
    sym_block3.initialize()
    a_data = mx.nd.ones((3,2))
    b_data = mx.nd.ones((3,2))
    sym_block3.optimize_for(a_data, b_data, backend='myProp')
    sym_filename, params_filename = sym_block3.export('optimized')
    assert sym_filename == 'optimized-symbol.json'
    assert params_filename is None

    # Test with additional input to subgraph op
    sym_block3.optimize_for(a_data, b_data, backend="addInputPass")
    out5 = sym_block3(a_data, b_data)

    # Reload exported block
    sym_block4 = nn.SymbolBlock.imports(sym_filename, ['a','b'], params_filename)

    out6 = sym_block4(a_data, b_data)
    # check that result matches one executed by MXNet
    assert_almost_equal(out[0].asnumpy(), out6[0].asnumpy(), rtol=1e-3, atol=1e-3)

@pytest.mark.skipif(check_platform(['x86_64']), reason="not all machine types supported")
@pytest.mark.skipif(is_cd_run(), reason="continuous delivery run - ignoring test")
def test_external_op():
    # check if operator already exists
    if hasattr(mx.nd, 'min_ex'):
        raise MXNetError('Operator already loaded')

    lib = 'libexternal_lib.so'
    fname = os.path.join(base_path,'example/extensions/lib_external_ops/build/'+lib)
    if not os.path.exists(fname):
        raise MXNetError(f"library {lib} not found ")

    fname = os.path.abspath(fname)
    mx.library.load(fname, False)

    # execute operator
    try:
        mx.nd.min_ex()
    except:
        raise MXNetError('Operator not loaded successfully')