# Copyright 2019 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.

"""operator dsl function: axpy"""

import akg.tvm
from akg.utils import validation_check as vc_util


def axpy(x, y, alpha):

    """
    Calculate the result of alpha * x + y

    The formula can be define as : \f$y=\alpha * x + y\f$

    Args:
        x (tvm.tensor.Tensor): Tensor of type float16, float32.
        y (tvm.tensor.Tensor): Tensor of type float16, float32.
        alpha (scalar): Affine transformation parameters.

    Returns:
        tvm.tensor.Tensor, has the same type and shape as x.
    """

    # Check types
    check_list = ["float16", "float32"]
    dtype = x.dtype
    if not dtype in check_list:
        raise RuntimeError("axpy_cce only support %s while dtype is %s" % (",".join(check_list), dtype))

    dtype = y.dtype
    if not dtype in check_list:
        raise RuntimeError("axpy_cce only support %s while dtype is %s" % (",".join(check_list), dtype))

    # Check shapes
    shape1 = [xx.value for xx in x.shape]
    shape2 = [yy.value for yy in y.shape]
    shapes = [shape1, shape2]
    for _, shape_val in enumerate(shapes):
        vc_util.check_shape(shape_val)

    # Calculate the result
    alpha = akg.tvm.const(alpha, dtype=dtype)
    res = akg.tvm.compute(shape1, lambda *indice: x(*indice) * alpha + y(*indice))
    return res
