# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed 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.
# ============================================================================
""" test_framstruct """
import numpy as np
import pytest

import mindspore.nn as nn
from mindspore import context
from mindspore.common import dtype as mstype
from mindspore.common.parameter import Parameter, ParameterTuple
from mindspore.common.tensor import Tensor
from mindspore.ops import composite as C
from mindspore.ops import operations as P
from mindspore.ops._grad.grad_base import bprop_getters
from mindspore.ops.primitive import prim_attr_register, PrimitiveWithInfer
from ..ut_filter import non_graph_engine
from ....mindspore_test_framework.utils.check_gradient import (
    ms_function, check_jacobian, Tensor, NNGradChecker,
    OperationGradChecker, check_gradient, ScalarGradChecker)


def setup_module(module):
    context.set_context(mode=context.PYNATIVE_MODE)


@ms_function
def while_upper_bound(upper):
    rval = 2
    while rval < upper:
        rval = rval * rval
    return rval


def test_while_upper_bound():
    res = while_upper_bound(10)
    assert res == 16


@ms_function
def while_lower_bound(lower):
    """ t_while """
    rval = lower
    while rval < 100:
        rval = rval * rval
    return rval


def test_while_lower_bound():
    res = while_lower_bound(2)
    assert res == 256


@ms_function
def dynamic_make_tuple(x, lower, upper):
    out = ()
    i = lower
    while i < upper:
        out = out + (x,)
        i = i + 1
    return out


def test_dynamic_make_tuple():
    # Dynamicly recursively creating static type is invalid in mindspore, as mindspore is a static language.
    with pytest.raises(RuntimeError):
        dynamic_make_tuple(2, 1, 5)


def test_make_tuple():
    # Staticly recursively creating static type is valid in mindspore.
    @ms_function
    def make_tuple(x):
        out = ()
        for i in range(3):
            out = out + (x,)
        return out

    res = make_tuple(5)
    assert res == (5, 5, 5)


@ms_function
def add(x, y):
    """ add """
    return x + y


def mul(x, y):
    """ mul """
    return x * y


def add_mul(x, y):
    """ add_mul """
    return (x + y) * y


def mainf(x, y):
    """ mainf """
    return C.grad_all(mul)(x, y)


def grad_add_mul(x, y):
    """ grad_add_mul """
    return C.grad_all(add_mul)(x, y)


@ms_function
def sub(x, y):
    """ sub """
    return x - y


@ms_function
def if_always_true(x):
    """ if_always_true """
    if True:
        return x
    else:
        return 0


def test_add():
    """ test_add """
    res = add(2.5, 3)
    assert res == 5.5


def test_sub():
    """ test_sub """
    res = sub(3.5, 3)
    assert res == 0.5


@non_graph_engine
def test_if_always_true():
    """ test_if_always_true """
    res = if_always_true(1)
    assert res == 1


@non_graph_engine
def test_f():
    """ test_f """
    res = mainf(3, 2)
    assert res == (2, 3)


@non_graph_engine
def test_grad_add_mul():
    """ test_grad_add_mul """
    res = grad_add_mul(3, 2)
    assert res == (2, 7)


def f(x):
    if x > 0:
        return f(x - 1)
    return x


@ms_function
def list_subscript():
    """ list_subscript """
    x = [1, 2, 3]
    return x[0] * x[1]


def test_list_subscript():
    """ test_list_subscript """
    res = list_subscript()
    assert res == 2


@ms_function
def ms_infer_for(xs, y):
    """ ms_infer_for """
    rval = y
    for x in xs:
        rval = rval + x
    return rval


def test_infer_for():
    """ test_infer_for """
    t = (1, 2, 3)
    y = 4
    res = ms_infer_for(t, y)
    assert res == 10


@ms_function
def if_construct(a, b):
    z = a
    if a > b:
        z = a + b
    else:
        z = a * b
    if z > b:
        return z - a
    else:
        return a - b


def test_if_construct():
    """ test_if_construct """
    res = if_construct(3, 6)
    assert res == 15


@ms_function
def if_scalar(a, b):
    """ if_abstract """
    if a:
        return a
    return b


def test_if_scalar1():
    """ test_if_abstract """
    res = if_scalar(3, 6)
    assert res == 3


def test_if_scalar2():
    """ test_if_abstract """
    res = if_scalar(0, 6)
    assert res == 6


@ms_function
def if_tensor(a, b):
    c = a
    if a < b:
        c = a + a
        if c < b:
            c = a + c
        else:
            c = a + b
    else:
        c = b + b
    out = c + c
    return out


def test_if_tensor():
    res = if_tensor(Tensor(np.ones([64, 10]).astype(np.int32)), Tensor(np.ones([64, 10]).astype(np.int32)))
    assert res == Tensor(np.ones([64, 10]).astype(np.int32) * 4)


@ms_function
def rec(x):
    """ rec """
    if x > 0:
        return rec(x - 1)
    return x


def test_grad_rec():
    """ test_grad_rec """
    res = C.grad(rec)(10)
    assert res == 1


def test_me_rec():
    """ test_me_rec """
    res = rec(10)
    assert res == 0


@ms_function
def t2_while(x, y):
    out = y - x
    i = 0
    while i < 10:
        out = mul(x, y)
        i = i + 1
    return out


def test_while2():
    res = t2_while(2, 3)
    assert res == 6


def test_grad_while2():
    res = C.grad(t2_while)(2, 3)
    assert res == 3


def if_test(a, b):
    """ if_test """
    if a > b:
        return 3 * a
    return 2 * b


def grad_if(x, y):
    """ grad_if """
    return C.grad_all(if_test)(x, y)


def test_grad_if():
    """ test_grad_if """
    assert grad_if(5, 4) == (3, 0)


# While loop is not unrolled in forward and backward graphs.
def test_dont_unroll_while():
    def dont_unroll_while(x, y):
        i = 2
        out = y - x
        while i < 10:
            out = mul(x, y)
            i = i + 1
        return out

    @ms_function()
    def invoke_while(x, y):
        return C.grad(dont_unroll_while)(x, y)

    res = invoke_while(2, 3)
    assert res == 3


class ConvNet(nn.Cell):
    def __init__(self):
        super(ConvNet, self).__init__()
        out_channel = 16
        kernel_size = 3
        self.conv = P.Conv2D(out_channel,
                             kernel_size,
                             mode=1,
                             pad_mode="pad",
                             pad=0,
                             stride=1,
                             dilation=2,
                             group=1)
        self.w = Parameter(Tensor(np.ones([16, 16, 3, 3]).astype(np.float32)), name='w')

    def construct(self, x):
        return self.conv(x, self.w)


conv = ConvNet()
c1 = Tensor([2], mstype.float32)
c2 = Tensor([10], mstype.float32)
c3 = Tensor([1], mstype.float32)


@ms_function
def t1_while(x, y, z):
    out = x
    i = c1
    while i < c2:
        out = out + conv(z)
        i = i + c3
    out = out + out
    return out


def test_while_net():
    y = Tensor(np.ones([1, 3, 3, 4]).astype(np.float32))
    x = Tensor(np.ones([1, 16, 12, 12]).astype(np.float32))
    z = Tensor(np.ones([1, 16, 16, 16]).astype(np.float32))
    res = t1_while(x, y, z)
    assert res == Tensor(np.ones([1, 16, 12, 12]).astype(np.float32) * 2306.0)


@ms_function
def if_while(a, b, x, z):
    c = a
    i = c1
    out = x
    if a < b:
        c = a + a
        while i < c2:
            out = out + conv(z)
            i = i + c3
    else:
        c = b + b
    out = c + c
    return out


def test_if_while():
    x = Tensor(np.random.randn(1, 16, 12, 12).astype(np.float32))
    z = Tensor(np.random.randn(1, 16, 16, 16).astype(np.float32))
    res = if_while(Tensor(np.ones([64, 10]).astype(np.float32)), Tensor(np.ones([64, 10]).astype(np.float32)), x, z)
    assert res == Tensor(np.ones([64, 10]).astype(np.float32) * 4.0)


def _while(x):
    """ _while """
    ret = x * x
    i = 2
    while i <= 3:
        ret = ret * i
        i = i + 1
    return ret


def grad_while(x):
    """ grad_while """
    return C.grad_all(_while)(x)


def test_grad_while():
    """ test_grad_while """
    assert grad_while(5) == (60,)


@ms_function
def factorial(n):
    """ factorial """
    if n == 0:
        return 1
    return n * factorial(n - 1)


def test_factorial():
    res = factorial(3)
    assert res == 6


def test_grad_factorial():
    res = C.grad(factorial)(3)
    assert res == 11


@ms_function
def factorial2(n):
    """ factorial """
    if n != 0:
        return n * factorial2(n - 1)
    elif n == 1:
        return 1 * factorial2(n - 1)
    else:
        return 1


def test_factorial2():
    res = factorial2(3)
    assert res == 6


@ms_function
def foo(n):
    if n <= 1:
        if n == 1:
            return foo(n - 1)
        else:
            return 1
    else:
        return foo(n - 1)


def test_foo():
    res = foo(5)
    assert res == 1


@ms_function
def double_nested_loop(x):
    i = 0
    s = 0
    while (i < x):
        j = 0
        i = i + 1
        while (j < 3):
            j = j + 1
            s = s + j
    return s


def test_nested_loop():
    res = double_nested_loop(3)
    assert res == 18


@ms_function
def double_nested_loop2(x):
    s = 0
    for i in range(x):
        for j in range(3):
            s = s + j
    return s


def test_nested_loop2():
    res = double_nested_loop(1)
    assert res == 6


def _for(x):
    """ _for """
    ret = x * x
    for i in (2, 3):
        ret = ret * i
    return ret


def grad_for(x):
    """ grad_for """
    return C.grad_all(_for)(x)


def test_grad_for():
    """ test_grad_for """
    assert grad_for(5) == (60,)


@ms_function
def try_tail(x):
    """ try_tail """
    return C.tail(x)


@non_graph_engine
def test_tail():
    """ test_tail """
    try_tail((0, 1, 2, 3))


@ms_function
def zero_like_tensor(x):
    """ zero_like_tensor """
    return C.zeros_like(x)


def test_zeros():
    """ test_zeros """
    x = Tensor(np.ones([2, 3]).astype(np.int32))
    res = zero_like_tensor(x)
    assert res == Tensor(np.zeros([2, 3]).astype(np.int32))


def test_ScalarGradChecker():
    """ test_ScalarGradChecker """

    def scalar_f(x, y):
        return x * y

    check_gradient(scalar_f, 1.0, 4.0, grad_checker_class=ScalarGradChecker, sampling_times=1)


def test_GradCheckerPrimitive():
    """ test_GradCheckerPrimitive """
    matmul = P.MatMul()

    def prim_f(x, y):
        return matmul(x, y)

    check_gradient(prim_f, Tensor(np.array([[0.65, 0.8, 0.8]], np.float32)),
                   Tensor(np.array([[0.1], [0.2], [-.1]], np.float32)),
                   grad_checker_class=OperationGradChecker, sampling_times=2)


def test_NNGradChecker():
    """ test_NNGradChecker """

    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()
            self.dense = nn.Dense(10, 10)

        def construct(self, x):
            out = self.dense(x)
            return out

    check_gradient(Net(), Tensor(np.random.rand(1, 10).astype(np.float32)),
                   delta=1e-3,
                   max_error=1e-3,
                   grad_checker_class=NNGradChecker, sampling_times=3)


def test_OperationGradChecker():
    """ test_OperationGradChecker """

    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()
            self.matmul = P.MatMul()
            self.z = Parameter(Tensor(np.array([1.0], np.float32)), name='z')

        def construct(self, x, y):
            x = x * self.z
            out = self.matmul(x, y)
            return out

    check_gradient(Net(), Tensor(np.array([[0.65, 0.8, 0.8]], np.float32)),
                   Tensor(np.array([[0.1], [0.2], [-.1]], np.float32)), grad_checker_class=OperationGradChecker,
                   input_selector=[1], sampling_times=2)


def test_ScalarJacobianChecker():
    """ test_ScalarJacobianChecker """

    def scalar_f(x, y):
        return x * y

    check_jacobian(scalar_f, 1.0, 4.0, grad_checker_class=ScalarGradChecker, input_selector=[0])


def test_OperationJacobianChecker():
    """ test_OperationJacobianChecker """

    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()
            self.matmul = P.MatMul()
            self.z = Parameter(Tensor(np.array([1.0], np.float32)), name='z')

        def construct(self, x, y):
            x = x * self.z
            out = self.matmul(x, y)
            return x, out

    check_jacobian(Net(), Tensor(np.array([[0.65, 0.8, 0.8], [0.1, 0.2, 0.3]], np.float32)),
                   Tensor(np.array([[0.1, 0.3], [0.2, 0.2], [-.1, 0.4]], np.float32)),
                   grad_checker_class=OperationGradChecker, input_selector=[0],
                   output_selector=[0])


def test_NNJacobianChecker():
    """ test_NNJacobianChecker """

    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()
            self.dense = nn.Dense(10, 10)

        def construct(self, x):
            out = self.dense(x)
            return out, x

    check_jacobian(Net(), Tensor(np.random.rand(1, 10).astype(np.float32)),
                   delta=1e-3,
                   max_error=1e-7,
                   grad_checker_class=NNGradChecker,
                   input_selector=[1],
                   output_selector=[0])


def multi_outputs(x, y):
    z = x + y
    return 2 * z, 2 * z


def test_grad_multi_outputs():
    assert C.grad_all_with_sens(multi_outputs)(2, 3, (1, 1)) == (4, 4)


@ms_function
def while_sp(x, y, z):
    out = x
    i = c3
    while i < c2:
        out = mul(x, out)
        i = i + c3
    return out


def test_while_sp():
    y = Tensor(np.ones([1, 3]).astype(np.float32))
    z = Tensor(np.ones([1, 3]).astype(np.float32))
    x = Tensor(np.ones([1, 3]).astype(np.float32) * 2.0)
    res = while_sp(x, y, z)
    assert res == Tensor(np.ones([1, 3]).astype(np.float32) * 1024.0)


def grad_refactor_simple_1(x, y):
    """ add """
    return x * x + 2 * y


def test_grad_refactor_simple_1():
    assert C.grad_all(grad_refactor_simple_1)(2, 1) == (4, 2)


def grad_refactor_simple_2(x, y, z):
    """ add """
    return x * y + z + x * y * z + x + x * y


def test_grad_refactor_simple_2():
    assert C.grad_all(grad_refactor_simple_2)(2, 3, 0) == (7, 4, 7)


def grad_refactor_1(a, b):
    """ if_test """

    def inner(x, y):
        return x * y

    return inner(a, b)


def test_grad_refactor_1():
    assert C.grad_all(grad_refactor_1)(2, 3) == (3, 2)


def grad_refactor_2(a, b):
    """ if_test """

    def inner(x):
        return x * b

    return inner(b) * inner(a)


def test_grad_refactor_2():
    assert C.grad_all(grad_refactor_2)(2, 3) == (27, 54)


def grad_refactor_3(a):
    """ if_test """
    if a > 3:
        return 0
    return 3 * a


def test_grad_refactor_3():
    assert C.grad_all(grad_refactor_3)(3) == (3,)


def grad_refactor_4(a):
    """ if_test """
    if a > 3:
        return 3 * a
    return 0


def test_grad_refactor_4():
    assert C.grad_all(grad_refactor_4)(4) == (3,)


def grad_refactor_5(a):
    """ if_test """
    if a > 3:
        return 1
    return a


def test_grad_refactor_5():
    assert C.grad_all(grad_refactor_5)(1) == (1,)


def grad_refactor_6(a, b):
    """ if_test """
    if a > b:
        return 3 * a + b
    return 2 * b * a


def test_grad_refactor_6():
    assert C.grad_all(grad_refactor_6)(3, 2) == (3, 1)


def grad_refactor_while(x):
    """ grad_refactor_while """
    rval = x
    while rval < 4:
        rval = rval * rval
    return rval


def test_grad_refactor_9():
    assert C.grad_all(grad_refactor_while)(3) == (6,)


def grad_refactor__while_1(x):
    """ _while """
    ret = x * x
    i = 2
    while i <= 3:
        ret = ret * i
        i = i + 1
    return ret


def test_grad_refactor_10():
    """ test_grad_while """
    assert C.grad_all(grad_refactor__while_1)(5) == (60,)


def test_grad_refactor_11():
    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()

        def construct(self, x, y):
            return x * y * y

    net = Net()
    C.grad_all(net)(Tensor(np.ones([2]).astype(np.float32)), Tensor(np.ones([2]).astype(np.float32)))


def test_grad_refactor_12():
    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()
            self.z = Parameter(Tensor(np.array([1.0], np.float32)), name='z')

        def construct(self, x, y):
            return x * self.z * y

    net = Net()
    C.grad_all(net)(Tensor(np.ones([2]).astype(np.float32)), Tensor(np.zeros([2]).astype(np.float32)))


def test_grad_refactor_13():
    class Net(nn.Cell):
        """ Net definition """

        def __init__(self):
            super(Net, self).__init__()
            self.z = Parameter(Tensor(np.ones([2]).astype(np.float32)), name='z')

        def construct(self, x, y):
            return x * self.z * y

    net = Net()
    weights = ParameterTuple(net.trainable_params())
    C.grad_by_list(net, weights)(Tensor(np.ones([2]).astype(np.float32)), Tensor(np.zeros([2]).astype(np.float32)))


def grad_refactor_14(a, b):
    """ if_test """

    def inner1(x):
        return x * b

    def inner2(x):
        return a * b

    def inner3(x):
        if (x > 2):
            return a
        return b

    return inner1(b) + inner2(a) + inner3(a)


def test_grad_refactor_14():
    assert C.grad_all(grad_refactor_14)(2, 3) == (3, 9)


class IfDeferInline(nn.Cell):
    def __init__(self, mul_size):
        super().__init__()
        self.mul_weight = Tensor(np.full(mul_size, 0.6, dtype=np.float32))
        self.mul = P.Mul()

    def construct(self, inputs):
        x = self.mul(inputs, self.mul_weight)
        if True:
            x = x
        return x


def test_grad_if_defer_inline():
    """ test_grad_if_defer_inline """
    network = IfDeferInline([128, 96])
    network.add_flags(defer_inline=False)
    inp = Tensor(np.ones([128, 96]).astype(np.float32))
    grads = C.grad_all(network)(inp)
    assert grads == (Tensor(np.full([128, 96], 0.6, dtype=np.float32)),)


def test_bprop_with_wrong_output_num():
    class BpropWithWrongOutputNum(PrimitiveWithInfer):
        @prim_attr_register
        def __init__(self):
            super(BpropWithWrongOutputNum, self).__init__('BpropWithWrongOutputNum')

        def __call__(self, x, y):
            return x

        def infer_shape(self, x_shape, yshape):
            return x_shape

        def infer_dtype(self, x_type, y_type):
            return x_type

    @bprop_getters.register(BpropWithWrongOutputNum)
    def get_bprop_with_wrong_output_num(self):
        """Generate bprop for BpropWithWrongOutputNum"""

        def bprop(x, y, out, dout):
            return (dout,)

        return bprop

    class BpropWithWrongOutputNumCell(nn.Cell):
        def __init__(self):
            super(BpropWithWrongOutputNumCell, self).__init__()

        def construct(self, x, y):
            return BpropWithWrongOutputNum()(x, y)

    with pytest.raises(TypeError):
        C.grad_all(BpropWithWrongOutputNumCell())(1, 2)


def test_bprop_with_wrong_output_type():
    class BpropWithWrongOutputType(PrimitiveWithInfer):
        @prim_attr_register
        def __init__(self):
            super(BpropWithWrongOutputType, self).__init__('BpropWithWrongOutputType')

        def __call__(self, x):
            return x

        def infer_shape(self, x_shape):
            return x_shape

        def infer_dtype(self, x_type):
            return x_type

    @bprop_getters.register(BpropWithWrongOutputType)
    def get_bprop_with_wrong_output_type(self):
        """Generate bprop for BpropWithWrongOutputType"""

        def bprop(x, out, dout):
            return (1,)

        return bprop

    class BpropWithWrongOutputTypeCell(nn.Cell):
        def __init__(self):
            super(BpropWithWrongOutputTypeCell, self).__init__()

        def construct(self, x):
            return BpropWithWrongOutputType()(x)

    with pytest.raises(TypeError):
        C.grad_all(BpropWithWrongOutputTypeCell())(Tensor(np.ones([64, 10]).astype(np.int32)))


def test_bprop_with_wrong_output_shape():
    class BpropWithWrongOutputShape(PrimitiveWithInfer):
        @prim_attr_register
        def __init__(self):
            super(BpropWithWrongOutputShape, self).__init__('BpropWithWrongOutputShape')

        def __call__(self, x):
            return x

        def infer_shape(self, x_shape):
            return x_shape

        def infer_dtype(self, x_type):
            return x_type

    @bprop_getters.register(BpropWithWrongOutputShape)
    def get_bprop_with_wrong_output_shape(self):
        """Generate bprop for BpropWithWrongOutputShape"""
        ones = Tensor(np.ones([2, ]).astype(np.int32))

        def bprop(x, out, dout):
            return (ones,)

        return bprop

    class BpropWithWrongOutputShapeCell(nn.Cell):
        def __init__(self):
            super(BpropWithWrongOutputShapeCell, self).__init__()

        def construct(self, x):
            return BpropWithWrongOutputShape()(x)

    with pytest.raises(TypeError):
        C.grad_all(BpropWithWrongOutputShapeCell())(Tensor(np.ones([64, 10]).astype(np.int32)))