# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

from __future__ import print_function

from ..framework import core
from ..fluid.layer_helper import LayerHelper
from ..fluid.data_feeder import check_variable_and_dtype

# TODO: define functions to get tensor attributes
from ..fluid.layers import rank # noqa: F401
from ..fluid.layers import shape # noqa: F401
import paddle
from paddle import _C_ops
from paddle.static import Variable
from ..fluid.framework import _in_legacy_dygraph, in_dygraph_mode

__all__ = []


def _complex_to_real_dtype(dtype):
 if dtype == core.VarDesc.VarType.COMPLEX64:
 return core.VarDesc.VarType.FP32
 elif dtype == core.VarDesc.VarType.COMPLEX128:
 return core.VarDesc.VarType.FP64
 else:
 return dtype


def _real_to_complex_dtype(dtype):
 if dtype == core.VarDesc.VarType.FP32:
 return core.VarDesc.VarType.COMPLEX64
 elif dtype == core.VarDesc.VarType.FP64:
 return core.VarDesc.VarType.COMPLEX128
 else:
 return dtype


def is_complex(x):
 """Return whether x is a tensor of complex data type(complex64 or complex128).

 Args:
 x (Tensor): The input tensor.

 Returns:
 bool: True if the data type of the input is complex data type, otherwise false.

 Examples:
 .. code-block:: python

 import paddle

 x = paddle.to_tensor([1 + 2j, 3 + 4j])
 print(paddle.is_complex(x))
 # True

 x = paddle.to_tensor([1.1, 1.2])
 print(paddle.is_complex(x))
 # False

 x = paddle.to_tensor([1, 2, 3])
 print(paddle.is_complex(x))
 # False
 """
 if not isinstance(x, (paddle.Tensor, paddle.static.Variable)):
 raise TypeError("Expected Tensor, but received type of x: {}".format(
 type(x)))
 dtype = x.dtype
 is_complex_dtype = (dtype == core.VarDesc.VarType.COMPLEX64 or
 dtype == core.VarDesc.VarType.COMPLEX128)
 return is_complex_dtype


def is_floating_point(x):
 """
 Returns whether the dtype of `x` is one of paddle.float64, paddle.float32, paddle.float16, and paddle.bfloat16.

 Args:
 x (Tensor): The input tensor.

 Returns:
 bool: True if the dtype of `x` is floating type, otherwise false.

 Examples:
 .. code-block:: python

 import paddle

 x = paddle.arange(1., 5., dtype='float32')
 y = paddle.arange(1, 5, dtype='int32')
 print(paddle.is_floating_point(x))
 # True
 print(paddle.is_floating_point(y))
 # False
 """
 if not isinstance(x, (paddle.Tensor, paddle.static.Variable)):
 raise TypeError("Expected Tensor, but received type of x: {}".format(
 type(x)))
 dtype = x.dtype
 is_fp_dtype = (dtype == core.VarDesc.VarType.FP32 or
 dtype == core.VarDesc.VarType.FP64 or
 dtype == core.VarDesc.VarType.FP16 or
 dtype == core.VarDesc.VarType.BF16)
 return is_fp_dtype


def is_integer(x):
 """Return whether x is a tensor of integeral data type.

 Args:
 x (Tensor): The input tensor.

 Returns:
 bool: True if the data type of the input is integer data type, otherwise false.

 Examples:
 .. code-block:: python

 import paddle

 x = paddle.to_tensor([1 + 2j, 3 + 4j])
 print(paddle.is_integer(x))
 # False

 x = paddle.to_tensor([1.1, 1.2])
 print(paddle.is_integer(x))
 # False

 x = paddle.to_tensor([1, 2, 3])
 print(paddle.is_integer(x))
 # True
 """
 if not isinstance(x, (paddle.Tensor, paddle.static.Variable)):
 raise TypeError("Expected Tensor, but received type of x: {}".format(
 type(x)))
 dtype = x.dtype
 is_int_dtype = (dtype == core.VarDesc.VarType.UINT8 or
 dtype == core.VarDesc.VarType.INT8 or
 dtype == core.VarDesc.VarType.INT16 or
 dtype == core.VarDesc.VarType.INT32 or
 dtype == core.VarDesc.VarType.INT64)
 return is_int_dtype


def real(x, name=None):
 """
 Returns a new tensor containing real values of the input tensor.

 Args:
 x (Tensor): the input tensor, its data type could be complex64 or complex128.
 name (str, optional): The default value is None. Normally there is no need for
 user to set this property. For more information, please refer to :ref:`api_guide_Name` .

 Returns:
 Tensor: a tensor containing real values of the input tensor.

 Examples:
 .. code-block:: python

 import paddle

 x = paddle.to_tensor(
 [[1 + 6j, 2 + 5j, 3 + 4j], [4 + 3j, 5 + 2j, 6 + 1j]])
 # Tensor(shape=[2, 3], dtype=complex64, place=CUDAPlace(0), stop_gradient=True,
 # [[(1+6j), (2+5j), (3+4j)],
 # [(4+3j), (5+2j), (6+1j)]])

 real_res = paddle.real(x)
 # Tensor(shape=[2, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
 # [[1., 2., 3.],
 # [4., 5., 6.]])

 real_t = x.real()
 # Tensor(shape=[2, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
 # [[1., 2., 3.],
 # [4., 5., 6.]])
 """
 if in_dygraph_mode():
 return _C_ops.final_state_real(x)
 if _in_legacy_dygraph():
 return _C_ops.real(x)

 check_variable_and_dtype(x, 'x', ['complex64', 'complex128'], 'real')
 helper = LayerHelper('real', **locals())
 out = helper.create_variable_for_type_inference(
 dtype=_complex_to_real_dtype(helper.input_dtype()))
 helper.append_op(type='real', inputs={'X': x}, outputs={'Out': out})
 return out


def imag(x, name=None):
 """
 Returns a new tensor containing imaginary values of input tensor.

 Args:
 x (Tensor): the input tensor, its data type could be complex64 or complex128.
 name (str, optional): The default value is None. Normally there is no need for
 user to set this property. For more information, please refer to :ref:`api_guide_Name` .

 Returns:
 Tensor: a tensor containing imaginary values of the input tensor.

 Examples:
 .. code-block:: python

 import paddle

 x = paddle.to_tensor(
 [[1 + 6j, 2 + 5j, 3 + 4j], [4 + 3j, 5 + 2j, 6 + 1j]])
 # Tensor(shape=[2, 3], dtype=complex64, place=CUDAPlace(0), stop_gradient=True,
 # [[(1+6j), (2+5j), (3+4j)],
 # [(4+3j), (5+2j), (6+1j)]])

 imag_res = paddle.imag(x)
 # Tensor(shape=[2, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
 # [[6., 5., 4.],
 # [3., 2., 1.]])

 imag_t = x.imag()
 # Tensor(shape=[2, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
 # [[6., 5., 4.],
 # [3., 2., 1.]])
 """
 if in_dygraph_mode():
 return _C_ops.final_state_imag(x)
 if _in_legacy_dygraph():
 return _C_ops.imag(x)

 check_variable_and_dtype(x, 'x', ['complex64', 'complex128'], 'imag')
 helper = LayerHelper('imag', **locals())
 out = helper.create_variable_for_type_inference(
 dtype=_complex_to_real_dtype(helper.input_dtype()))
 helper.append_op(type='imag', inputs={'X': x}, outputs={'Out': out})
 return out