# 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 control ops """
import numpy as np
import mindspore as ms
from mindspore import nn
from mindspore import Tensor
from mindspore import context
from mindspore.ops import operations as P
from mindspore.ops import composite as C
from mindspore.ops import functional as F
from mindspore.common.parameter import Parameter, ParameterTuple

context.set_context(mode=context.GRAPH_MODE)


def cond_data_test(x_init, y_init):
    class Net(nn.Cell):
        def __init__(self):
            """"""
            super(Net, self).__init__()
            self.square = P.Square()
            self.add = P.TensorAdd()
            self.value = Tensor(np.full((1), 3, dtype=np.float32))
            self.switch = P.GeSwitch()
            self.merge = P.Merge()
            self.less = P.Less()

        def construct(self, x, y):
            cond = self.less(x, y)
            st1, sf1 = self.switch(x, cond)
            st2, sf2 = self.switch(y, cond)
            add_ret = self.add(st1, st2)
            st3, sf3 = self.switch(self.value, cond)
            sq_ret = self.square(sf3)
            ret = self.merge((add_ret, sq_ret))
            return ret[0]

    x = Tensor(x_init, dtype=ms.float32)
    y = Tensor(y_init, dtype=ms.float32)
    net = Net()
    output = net(x, y)
    return output


def test_cond_data_true():
    output = cond_data_test(3, 8)
    print("test_cond_data_true:", output)

def test_cond_data_false():
    output = cond_data_test(8, 3)
    print("test_cond_data_false:", output)

def if_compile_test(x_init, y_init):
    class Net(nn.Cell):
        def __init__(self):
            """"""
            super(Net, self).__init__()
            self.square = P.Square()
            self.add = P.TensorAdd()
            self.value = Tensor(3, dtype=ms.float32)
            self.switch = P.GeSwitch()
            self.merge = P.Merge()
            self.less = P.Less()

        def construct(self, x, y):
            cond = self.less(x, y)
            ret = self.value
            if cond:
                ret = self.add(x, ret)
                ret = self.add(y, ret)
            else:
                ret = self.square(self.value)
            return ret

    x = Tensor(x_init, dtype=ms.float32)
    y = Tensor(y_init, dtype=ms.float32)
    net = Net()
    output = net(x, y)
    return output


def test_if_none():
    class Net(nn.Cell):
        def __init__(self, z: None):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = None
    net = Net(z)
    assert net(x, y) == y


def test_if_str_is_not_none_right():
    class Net(nn.Cell):
        def __init__(self, z: str):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z == None:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = "ok"
    net = Net(z)
    assert net(x, y) == y


def test_if_str_is_not_none_left():
    class Net(nn.Cell):
        def __init__(self, z: str):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if None == self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = "ok"
    net = Net(z)
    assert net(x, y) == y


def test_if_none_equal_none():
    class Net(nn.Cell):
        def __init__(self, z: None):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z == None:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = None
    net = Net(z)
    assert net(x, y) == x


def test_if_str_is_null():
    class Net(nn.Cell):
        def __init__(self, z: str):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = ""
    net = Net(z)
    assert net(x, y) == y


def test_if_str_is_true():
    class Net(nn.Cell):
        def __init__(self, z: str):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 9, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = "ok"
    net = Net(z)
    assert net(x, y) == x


def test_if_str_equal():
    class Net(nn.Cell):
        def __init__(self, z: str):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z == "ok":
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = "ok"
    net = Net(z)
    assert net(x, y) == x


def test_if_tuple_is_null():
    class Net(nn.Cell):
        def __init__(self, z: tuple):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = ()
    net = Net(z)
    assert net(x, y) == y


def test_if_tuple_is_not_null():
    class Net(nn.Cell):
        def __init__(self, z: tuple):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = (1, 2, 3)
    net = Net(z)
    assert net(x, y) == x


def test_if_dict_is_null():
    class Net(nn.Cell):
        def __init__(self, z: dict):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = {}
    net = Net(z)
    assert net(x, y) == y


def test_if_dict_is_not_null():
    class Net(nn.Cell):
        def __init__(self, z: dict):
            """"""
            super(Net, self).__init__()
            self.z = z

        def construct(self, x, y):
            if self.z:
                ret = x
            else:
                ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = {"one": 1, "two": 2}
    net = Net(z)
    assert net(x, y) == x


def test_if_else_assign():
    class Net(nn.Cell):
        def __init__(self, m: list):
            """"""
            super(Net, self).__init__()
            self.m = m
            self.n = [4, 5, 6]

        def construct(self, x, y):
            exp_1 = self.m if self.m else self.n
            exp_2 = self.m if exp_1 == self.n else self.n
            if exp_2 == self.m:
                if self.m:
                    ret = x
                else:
                    ret = y
            else:
                if self.m:
                    ret = x
                else:
                    ret = y
            return ret

    x = Tensor(np.ones([6, 8, 10], np.int32))
    y = Tensor(np.zeros([3, 4, 5], np.int32))
    z = [1, 2]
    net = Net(z)
    assert net(x, y) == x


def test_if_compile_true():
    output = if_compile_test(3, 8)
    print("test_if_compile_true:", output)


def test_if_compile_false():
    output = if_compile_test(8, 3)
    print("test_if_compile_false:", output)


def test_switch_layer():
    class Layer1(nn.Cell):
        def __init__(self):
            super(Layer1, self).__init__()
            self.z1 = Parameter(Tensor(np.full([128, 96], 0.6, dtype=np.float32)), name='z1')
        def construct(self, x):
            return x * self.z1

    class Layer2(nn.Cell):
        def __init__(self):
            super(Layer2, self).__init__()
            self.z2 = Parameter(Tensor(np.full([128, 96], 0.6, dtype=np.float32)), name='z2')
        def construct(self, x):
            return x * self.z2

    class SwitchLayerCell(nn.Cell):
        def __init__(self):
            super(SwitchLayerCell, self).__init__()
            self.layers = (Layer1(), Layer2())
            self.z3 = Parameter(Tensor(np.full([128, 96], 0.6, dtype=np.float32)), name='z3')
        def construct(self, index, x):
            ret = F.switch_layer(index, self.layers)(x) * self.z3
            return ret

    net = SwitchLayerCell()
    net(1, Tensor(np.full([128, 96], 0.6, dtype=np.float32)))
    C.grad_by_list(net, ParameterTuple(net.trainable_params()))(0, Tensor(np.full([128, 96], 0.6, dtype=np.float32)))
    C.grad_all(net)(0, Tensor(np.full([128, 96], 0.6, dtype=np.float32)))