```# Copyright 2022 Huawei Technologies Co., Ltd
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# ============================================================================
from mindspore.common import Tensor
from mindspore import ops
from mindspore.ops.primitive import constexpr
from ..cell import Cell

@constexpr
r"""
Note: the input 'padding' in this function is already converted to list of lists to match MirrorPad
"""
raise ValueError(f"For padding with length {len(padding) * 2}, the dimension of the tensor should be at least "

r"""
Check whether the input padding is a tuple or a int converted to a tuple.
Check if the length of padding is divisible by 2.
"""

@constexpr
"""
Check relationship between input shape and padding to make sure after negative dimension padding the out is
positive.
"""
msg = "For '{}', the dimension of input must more than or equal to len(padding)/2, " \
"but got {}".format(name, len(input_shape))
raise ValueError(msg)
padding = [(0, 0) for i in range(len(input_shape) - 1)] + [padding]
else:
if index == 0:
dim_name = '1st'
elif index == 1:
dim_name = '2nd'
elif index == 2:
dim_name = '3rd'
else:
dim_name = str(index + 1) + 'th'

if item[0] < -input_shape[index]:
msg = "For '{}', the shape of input after padding must be positive, the input shape is {}, " \
"value of parameter 'padding' applied to the {} dimension of input must " \
"no less than -{}, but got {}".format(name, input_shape, dim_name, input_shape[index], item[0])
raise ValueError(msg)
if item[1] < -input_shape[index]:
msg = "For '{}', the shape of input after padding must be positive, the input shape is {}, " \
"value of parameter 'padding' applied to the {} dimension of input must " \
"no less than -{}, but got {}".format(name, input_shape, dim_name, input_shape[index], item[0])
raise ValueError(msg)
if input_shape[index] + item[0] + item[1] <= 0:
msg = "For '{}', the shape of input after padding must be positive, the input shape is {}, " \
"but the {} dimension of input shape {} plus padding {} and {} resulted in a non-positive output " \
"shape.".format(name, input_shape, dim_name, input_shape[index], item[0], item[1])
raise ValueError(msg)

@constexpr
"""get non-negative padding and make negative position."""
start = [0 for i in range(len(new_padding))]
end = [0 for i in range(len(new_padding))]
if item[0] < 0:
start[index] = item[0]
if item[1] < 0:
end[index] = item[1]

@constexpr
def _get_begin_size(shape, begin, end):
"""Calculate begin and size for ops.Slice."""
size = tuple([shape[i] + begin[i] + end[i] for i in range(len(shape))])
begin = tuple([int(-i) for i in begin])
return begin, size

r"""
Using a given value to pads the last n dimensions of input tensor.

Args:
sequence is starting from the last dimension and moving forward. The length of padding must be

The padding sequence is starting from the last dimension and moving forward.
the size of last dimension of output is :math:`padding\_0 + x.shape[-1] + padding\_1`.
The size of penultimate dimension of output is :math:`padding\_2 + x.shape[-2] + padding\_3`.
The size of 3rd to last dimension of output is :math:`padding\_4 + x.shape[-3] + padding\_5`.
The size of i-td to last dimension of output is :math:`padding\_{2m} + x.shape[-m-1] + padding\_{2m+1}`.
The remaining dimensions of the output are consistent with those of the input.
name (str): Name of method, used for positioning error messages in the base class.

Returns:

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Raises:
ValueError: If the length of padding is not a multiple of 2.
ValueError: If the length of input less than len(padding)/2.
ValueError: If the output shape after padding is not positive.
"""

if len(padding) % 2 != 0:
msg = "For '{}', the length of parameter 'padding' with tuple type must be a multiple of 2, " \
raise ValueError(msg)
msg = "For '{}', the length of parameter 'padding' with tuple type must equal to 2." \
raise ValueError(msg)
msg = "For '{}', the length of parameter 'padding' with tuple type must no more than 4." \
raise ValueError(msg)
msg = "For '{}', the length of parameter 'padding' with tuple type must no more than 6." \
raise ValueError(msg)

else:
msg = "For '{}', the type of parameter 'padding' must be in [int, float], " \
raise TypeError(msg)

if not isinstance(value, (int, float)):
msg = "For '{}', the type of parameter 'value' must be in [int, float], " \
"but got {}".format(name, type(value))
raise TypeError(msg)

self.value = value
self._name = name

def construct(self, x):
input_shape = x.shape
input_type = x.dtype
ones = ops.Ones()(output.shape, output.dtype)
value = ops.Fill()(output.dtype, output.shape, self.value)
slice_op = ops.Slice()
begin, size = _get_begin_size(output.shape, start, end)
output = slice_op(output, begin, size)
return output

r"""
Using a given constant value to pads the last dimensions of input tensor.

Args:
If is int, uses the same padding in both boundaries of input's last dimension.
If a 2-tuple, uses (padding_0, padding_1) to pad. If the input is `x`, the size of last
dimension of output is :math:`padding\_0 + x.shape[-1] + padding\_1`. The remaining dimensions
of the output are consistent with those of the input.

Returns:

Raises:
TypeError: If `padding` is not a tuple or int.
TypeError: If `value` is not int or float.
ValueError: If the length of `padding` with tuple type is not equal to 2.
ValueError: If the output shape after padding is not positive.

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Examples:
>>> import numpy as np
>>> from mindspore import Tensor
>>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
>>> x = Tensor(x)
>>> value = 0.5
>>> print(out)
[[[[1.  1.  1.  1.  0.5]
[1.  1.  1.  1.  0.5]
[1.  1.  1.  1.  0.5]]
[[1.  1.  1.  1.  0.5]
[1.  1.  1.  1.  0.5]
[1.  1.  1.  1.  0.5]]]]
>>> print(out.shape)
(1, 2, 3, 5)
>>> value = 0.5
>>> print(out)
[[[[0.5 1.  1.  1.  1.  0.5]
[0.5 1.  1.  1.  1.  0.5]
[0.5 1.  1.  1.  1.  0.5]]
[[0.5 1.  1.  1.  1.  0.5]
[0.5 1.  1.  1.  1.  0.5]
[0.5 1.  1.  1.  1.  0.5]]]]
>>> print(out.shape)
(1, 2, 3, 6)
>>> value = 0.5
>>> print(out)
[[[[1. 1. 1.]
[1. 1. 1.]
[1. 1. 1.]]
[[1. 1. 1.]
[1. 1. 1.]
[1. 1. 1.]]]]
>>> print(out.shape)
(1, 2, 3, 3)
"""

r"""
Using a given constant value to pads the last two dimensions of input tensor.

Args:
If is int, uses the same padding in boundaries of input's last two dimensions.
If the input is `x`, the size of last dimension of output is :math:`padding\_0 + x.shape[-1] + padding\_1`.
The size of penultimate dimension of output is :math:`padding\_2 + x.shape[-2] + padding\_3`.
The remaining dimensions of the output are consistent with those of the input.

Returns:

Raises:
TypeError: If `padding` is not a tuple or int.
TypeError: If `value` is not int or float.
ValueError: If the length of `padding` is more than 4 or not a multiple of 2.
ValueError: If the output shape after padding is not positive.

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Examples:
>>> import numpy as np
>>> from mindspore import Tensor
>>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
>>> x = Tensor(x)
>>> padding = (-1, 1, 0, 1)
>>> value = 0.5
>>> print(out)
[[[[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[0.5 0.5 0.5 0.5]]
[[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[0.5 0.5 0.5 0.5]]]]
>>> print(out.shape)
(1, 2, 4, 4)
"""

r"""
Using a given constant value to pads the last three dimensions of input tensor.

Args:
If is int, uses the same padding in boundaries of input's last three dimensions.
If is tuple and length of padding is 6 uses
the size of last dimension of output is :math:`padding\_0 + x.shape[-1] + padding\_1`.
The size of penultimate dimension of output is :math:`padding\_2 + x.shape[-2] + padding\_3`.
The size of 3rd to last dimension of output is :math:`padding\_4 + x.shape[-3] + padding\_5`.
The remaining dimensions of the output are consistent with those of the input.

Returns:

Raises:
TypeError: If `padding` is not a tuple or int.
TypeError: If `value` is not int or float.
ValueError: If the length of `padding` is more than 6 or not a multiple of 2.
ValueError: If the output shape after padding is not positive.

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Examples:
>>> import numpy as np
>>> from mindspore import Tensor
>>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
>>> x = Tensor(x)
>>> padding = (-1, 1, 0, 1, 1, 0)
>>> value = 0.5
>>> print(out)
[[[[0.5 0.5 0.5 0.5]
[0.5 0.5 0.5 0.5]
[0.5 0.5 0.5 0.5]
[0.5 0.5 0.5 0.5]]
[[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[0.5 0.5 0.5 0.5]]
[[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[1.  1.  1.  0.5]
[0.5 0.5 0.5 0.5]]]]
>>> print(out.shape)
(1, 3, 4, 4)
"""

r"""
Using a given padding to do reflection pad on the given tensor.
Work as a parent class, and only accepts tuple as padding input.
"""

self.name = name
# check if padding and its elements are valid
raise TypeError(f"For '{self.name}' the input 'padding' must be an integer or tuple, "
if len(padding) % 2 != 0:
raise ValueError(f"For '{self.name}' the length of input 'padding' must be divisible by 2, "
if not all(isinstance(i, int) for i in padding):
raise TypeError(f"For '{self.name}' every element in 'padding' must be integer, "
if not all(i >= 0 for i in padding):
raise ValueError(f"For '{self.name}' every element in 'padding' must be >= 0. "

def construct(self, x):
input_shape = x.shape
return x

r"""
Using a given padding to do reflection pad on the given tensor.

Args:
If padding is an integer: all directions will be padded with the same size.

Inputs:
- **x** (Tensor) - 2D or 3D, shape: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.

Outputs:
Tensor, after padding. Shape: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`,

Raises:
TypeError: If 'padding' is not a tuple or int.
TypeError: If there is an element in 'padding' that is not int.
ValueError: If the length of 'padding' is not divisible by 2.
ValueError: If there is an element in 'padding' that is negative.
ValueError: If the there is a dimension mismatch between the padding and the tensor.

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Examples:
>>> import numpy as np
>>> from mindspore import Tensor
>>> x = Tensor(np.array([[[0, 1, 2, 3], [4, 5, 6, 7]]]).astype(np.float32))
>>> # x has shape (1, 2, 4)
>>> # The first and the second dimension of x remain the same.
>>> # The third dimension of x: W_out = W_in + pad_left + pad_right = 4 + 3 + 1 = 8
>>> # The shape of out is (1, 2, 8)
>>> print(out)
[[[3. 2. 1. 0. 1. 2. 3. 2.]
[7. 6. 5. 4. 5. 6. 7. 6.]]]
"""

r"""

Args:
If padding is an integer: all directions will be padded with the same size.

Inputs:
- **x** (Tensor) - 3D or 4D, shape: :math:`(C, H_{in}, W_{out})` or :math:`(N, C, H_{out}, W_{out})`.

Outputs:
Tensor, after padding. Shape: :math:`(C, H_{out}, W_{out})` or :math:`(N, C, H_{out}, W_{out})`,

Raises:
TypeError: If 'padding' is not a tuple or int.
TypeError: If there is an element in 'padding' that is not int.
ValueError: If the length of 'padding' is not divisible by 2.
ValueError: If there is an element in 'padding' that is negative.
ValueError: If the there is a dimension mismatch between the padding and the tensor.

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Examples:
>>> import numpy as np
>>> from mindspore import Tensor
>>> x = Tensor(np.array([[[0, 1, 2], [3, 4, 5], [6, 7, 8]]]).astype(np.float32))
>>> # x has shape (1, 3, 3)
>>> padding = (1, 1, 2, 0)
>>> # The first dimension of x remains the same.
>>> # The second dimension of x: H_out = H_in + pad_up + pad_down = 3 + 1 + 1 = 5
>>> # The third dimension of x: W_out = W_in + pad_left + pad_right = 3 + 2 + 0 = 5
>>> # The shape of out is (1, 5, 5)
>>> print(out)
[[[7. 6. 7. 8. 7.]
[4. 3. 4. 5. 4.]
[1. 0. 1. 2. 1.]
[4. 3. 4. 5. 4.]
[7. 6. 7. 8. 7.]]]
"""

r"""
Pads the last two dimensions of input tensor with zero.

Args:
If is int, uses the same padding in boundaries of input's last two dimensions.
If the input is `x`, the size of last dimension of output is :math:`padding\_0 + x.shape[-1] + padding\_1`.
The size of penultimate dimension of output is :math:`padding\_2 + x.shape[-2] + padding\_3`.
The remaining dimensions of the output are consistent with those of the input.

Returns:

Raises:
TypeError: If `padding` is not a tuple or int.
ValueError: If the length of `padding` is more than 4 or not a multiple of 2.
ValueError: If the output shape after padding is not positive.

Supported Platforms:
``Ascend`` ``GPU`` ``CPU``

Examples:
>>> import numpy as np
>>> from mindspore import Tensor
>>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
>>> x = Tensor(x)
>>> padding = (-1, 1, 0, 1)
>>> print(out)
[[[[1. 1. 1. 0.]
[1. 1. 1. 0.]
[1. 1. 1. 0.]
[0. 0. 0. 0.]]
[[1. 1. 1. 0.]
[1. 1. 1. 0.]
[1. 1. 1. 0.]
[0. 0. 0. 0.]]]]
>>> print(out.shape)
(1, 2, 4, 4)
"""