# 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.
# ============================================================================
"""
@File  : test_tensor.py
@Author:
@Date  : 2019-03-14
@Desc  : test mindspore tensor's operation
"""
import numpy as np
import pytest

import mindspore as ms
import mindspore.common.api as me
import mindspore.nn as nn
from ..ut_filter import non_graph_engine


ndarr = np.ones((2, 3))

def test_tensor_flatten():
    with pytest.raises(AttributeError):
        lst = [1, 2, 3, 4,]
        tensor_list = ms.Tensor(lst, ms.float32)
        tensor_list = tensor_list.Flatten()
        print(tensor_list)

def test_tensor_list():
    lst = [[1.0, 2.0, 1.0], [1.0, 10.0, 9.0]]
    tensor_list = ms.Tensor(lst, ms.float32)
    print(tensor_list)

def test_tensor():
    """test_tensor"""
    t1 = ms.Tensor(ndarr)
    assert isinstance(t1, ms.Tensor)
    assert t1.dtype() == ms.float64

    t2 = ms.Tensor(np.zeros([1, 2, 3]), ms.float32)
    assert isinstance(t2, ms.Tensor)
    assert t2.shape() == (1, 2, 3)
    assert t2.dtype() == ms.float32

    t3 = ms.Tensor(0.1)
    assert isinstance(t3, ms.Tensor)
    assert t3.dtype() == ms.float64

    t4 = ms.Tensor(1)
    assert isinstance(t4, ms.Tensor)
    assert t4.dtype() == ms.int64

def test_tensor_type_float16():
    t_float16 = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16))
    assert isinstance(t_float16, ms.Tensor)
    assert t_float16.shape() == (2, 3)
    assert t_float16.dtype() == ms.float16


def test_tensor_type_float32():
    t_float32 = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32))
    assert isinstance(t_float32, ms.Tensor)
    assert t_float32.shape() == (2, 3)
    assert t_float32.dtype() == ms.float32


def test_tensor_type_float32_user_define():
    t = ms.Tensor(np.zeros([1, 2, 3]), ms.float32)
    assert isinstance(t, ms.Tensor)
    assert t.shape() == (1, 2, 3)
    assert t.dtype() == ms.float32


def test_tensor_type_float64():
    t = ms.Tensor([[1.0, 2, 3], [4, 5, 6]])
    assert isinstance(t, ms.Tensor)
    assert t.shape() == (2, 3)
    assert t.dtype() == ms.float64

    t_zero = ms.Tensor(np.zeros([1, 2, 3]))
    assert isinstance(t_zero, ms.Tensor)
    assert t_zero.shape() == (1, 2, 3)
    assert t_zero.dtype() == ms.float64


def test_tensor_type_float64_user_define():
    t = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=float))
    assert isinstance(t, ms.Tensor)
    assert t.shape() == (2, 3)
    assert t.dtype() == ms.float64

    t_float64 = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]]), ms.float64)
    assert isinstance(t_float64, ms.Tensor)
    assert t_float64.shape() == (2, 3)
    assert t_float64.dtype() == ms.float64

def test_tensor_type_bool():
    # init a tensor with bool type
    ts_bool_array = ms.Tensor(np.zeros([2, 3], np.bool), ms.bool_)
    assert isinstance(ts_bool_array, ms.Tensor)
    assert ts_bool_array.dtype() == ms.bool_

    t_bool = ms.Tensor(True)
    assert isinstance(t_bool, ms.Tensor)
    assert t_bool.dtype() == ms.bool_

    t_bool_array = ms.Tensor(np.array([[True, False, True], [False, False, False]]))
    assert isinstance(t_bool_array, ms.Tensor)
    assert t_bool_array.shape() == (2, 3)
    assert t_bool_array.dtype() == ms.bool_

def test_tensor_type_int8():
    t_int8_array = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int8))
    assert isinstance(t_int8_array, ms.Tensor)
    assert t_int8_array.shape() == (2, 3)
    assert t_int8_array.dtype() == ms.int8


def test_tensor_type_int16():
    t_int16_array = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int16))
    assert isinstance(t_int16_array, ms.Tensor)
    assert t_int16_array.shape() == (2, 3)
    assert t_int16_array.dtype() == ms.int16


def test_tensor_type_int32():
    t_int = ms.Tensor([[1, 2, 3], [4, 5, 6]])
    assert isinstance(t_int, ms.Tensor)
    assert t_int.shape() == (2, 3)
    assert t_int.dtype() == ms.int64

    t_int_array = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32))
    assert isinstance(t_int_array, ms.Tensor)
    assert t_int_array.shape() == (2, 3)
    assert t_int_array.dtype() == ms.int32


def test_tensor_type_int64():
    t_int64 = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64))
    assert isinstance(t_int64, ms.Tensor)
    assert t_int64.shape() == (2, 3)
    assert t_int64.dtype() == ms.int64

def test_tensor_type_uint8():
    t_uint8_array = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8))
    assert isinstance(t_uint8_array, ms.Tensor)
    assert t_uint8_array.shape() == (2, 3)
    assert t_uint8_array.dtype() == ms.uint8


def test_tensor_type_uint16():
    t_uint16_array = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint16))
    assert isinstance(t_uint16_array, ms.Tensor)
    assert t_uint16_array.shape() == (2, 3)
    assert t_uint16_array.dtype() == ms.uint16


def test_tensor_type_uint32():
    t_uint32_array = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32))
    assert isinstance(t_uint32_array, ms.Tensor)
    assert t_uint32_array.shape() == (2, 3)
    assert t_uint32_array.dtype() == ms.uint32


def test_tensor_type_uint64():
    t_uint64 = ms.Tensor(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64))
    assert isinstance(t_uint64, ms.Tensor)
    assert t_uint64.shape() == (2, 3)
    assert t_uint64.dtype() == ms.uint64

def test_set_type():
    t = ms.Tensor(ndarr)
    t.set_dtype(ms.float32)
    assert t.dtype() == ms.float32


@non_graph_engine
def test_add():
    x = ms.Tensor(ndarr)
    y = ms.Tensor(ndarr)
    z = x + y
    assert isinstance(z, ms.Tensor)


@non_graph_engine
def test_sub():
    x = ms.Tensor(ndarr)
    y = ms.Tensor(ndarr)
    z = x - y
    assert isinstance(z, ms.Tensor)


class Net(nn.Cell):
    """Net definition"""
    def __init__(self, dim):
        super(Net, self).__init__()
        self.dim = dim

    def construct(self, input_x):
        return input_x


@non_graph_engine
def test_return_tensor():
    """test_return_tensor"""
    net = Net(0)
    input_data = ms.Tensor(np.array([[1.2, 2.1], [2.2, 3.2]]).astype('float32'))
    input_data.set_dtype(ms.float32)
    exe = me._executor
    exe.compile(net, input_data)
    tensor_ = exe(net, input_data)

    # get shape
    shape_ = tensor_.shape()
    print("shape = ", shape_)

    # get type
    type_ = tensor_.dtype()
    print("type = ", type_)

    # get value
    value_ = tensor_.asnumpy()
    print("numpy value = ", value_)


def test_tensor_contiguous():
    """test_tensor_contiguous"""
    input_c = np.arange(6).reshape(2, 3)
    input_f = input_c.T
    np.ascontiguousarray(input_c, dtype=np.float32)
    assert True, input_c.flags['C_CONTIGUOUS']

    print("input_f flags = ", input_f.flags)
    assert True, input_f.flags['F_CONTIGUOUS']

    tensor_f_float32 = ms.Tensor(input_f)
    rt_f = tensor_f_float32.asnumpy()
    assert True, rt_f.flags['C_CONTIGUOUS']
    print("rt_f flags = ", rt_f.flags)

def test_tensor_contiguous2():
    input_data = np.random.randn(32, 112, 112, 3).astype(np.float32)
    input_me = input_data.transpose(0, 3, 1, 2)
    print("input_me flags = ", input_me.flags)
    tensor_f_float32 = ms.Tensor(input_me)
    out_f = tensor_f_float32.asnumpy()
    print("out_f flags = ", out_f.flags)

def test_tensor_input_string():
    with pytest.raises(TypeError):
        input_data = 'ccc'
        ms.Tensor(input_data)

def test_tensor_input_tuple_string():
    with pytest.raises(TypeError):
        input_data = (2, 3, '4', 5)
        ms.Tensor(input_data)

def test_tensor_input_list_string():
    with pytest.raises(TypeError):
        input_data = [[2, 3, '4', 5], [1, 2, 3, 4]]
        ms.Tensor(input_data)

def test_tensor_input_none():
    with pytest.raises(TypeError):
        input_data = None
        ms.Tensor(input_data, np.int64)

# pylint: disable=no-value-for-parameter
def test_tensor_input_empty():
    with pytest.raises(TypeError):
        ms.Tensor()

def test_tensor_input_ndarray_str():
    with pytest.raises(TypeError):
        inp = np.array(["88", 2, 4])
        ms.Tensor(inp)

def test_tensor_input_ndarray_bool():
    inp = np.array([True, 2, 4])
    ms.Tensor(inp)

    inp = np.array([False, 2, 4])
    ms.Tensor(inp)

def test_tensor_input_ndarray_complex():
    with pytest.raises(TypeError):
        inp = np.array([20j, 2, 4])
        ms.Tensor(inp)

def test_tensor_input_ndarray_none():
    with pytest.raises(TypeError):
        inp = np.array([None, 2, 4])
        ms.Tensor(inp)

def test_tensor_input_ndarray_dict():
    with pytest.raises(TypeError):
        inp = {'a': 6, 'b': 7}
        ms.Tensor(inp)

def test_tensor_input_np_nan():
    with pytest.raises(TypeError):
        input_data = (1, 2, 3, np.nan)
        ms.Tensor(input_data, np.int64)

def test_tensor_input_tuple_inf():
    with pytest.raises(TypeError):
        input_data = (1, 2, 3, float("inf"))
        ms.Tensor(input_data, np.int64)

def test_tensor_input_dict():
    with pytest.raises(TypeError):
        input_data = {'a': 6, 'b': 7}
        ms.Tensor(input_data, np.int64)

def test_tensor_input_complex():
    with pytest.raises(TypeError):
        input_data = (1, 2j, 3)
        ms.Tensor(input_data, np.int64)

def test_tensor_dtype_np_float():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.float)
        ms.Tensor(input_data, np.float)

def test_tensor_dtype_np_float16():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.float16)
        ms.Tensor(input_data, np.float16)

def test_tensor_dtype_np_float32():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.float32)
        ms.Tensor(input_data, np.float32)

def test_tensor_dtype_np_float64():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.float64)
        ms.Tensor(input_data, np.float64)

def test_tensor_dtype_np_int():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.int)
        ms.Tensor(input_data, np.int)

def test_tensor_dtype_np_int8():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.int8)
        ms.Tensor(input_data, np.int8)

def test_tensor_dtype_np_int16():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.int16)
        ms.Tensor(input_data, np.int16)

def test_tensor_dtype_np_int32():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.int32)
        ms.Tensor(input_data, np.int32)

def test_tensor_dtype_np_int64():
    with pytest.raises(TypeError):
        input_data = np.random.randn(32, 112, 112, 3).astype(np.int64)
        ms.Tensor(input_data, np.int64)

def test_tensor_dtype_fp32_to_bool():
    with pytest.raises(RuntimeError):
        input = np.random.randn(2, 3, 4, 5).astype(np.float32)
        input = ms.Tensor(input)
        input_me = ms.Tensor(input, dtype=ms.bool_)