# SPDX-License-Identifier: Apache-2.0

from .....proto import onnx_proto
from ....common._apply_operation import apply_add, apply_mul
from ....common._registration import register_converter


def deduce_broadcast_axis_and_shape(target_opset, shape):
    # This function is used to calculate the first axis aligned with the scalar and the scalar's ONNX shape for reduce-
    # like operators. Assuming input variable is always a 4-D tensor, we provide a few of examples. If scalar's shape
    # is [1, 2, 3] and input shape is [5, 2, 3, 8], the aligned axis is the [2] (indexed by 1 because indexes are 0-based)
    # in [5, 2, 3, 8], and the desired scalar shape in ONNX is [2, 3] # (the [1] in [1, 2, 3] is redundant and can cause
    # errors in ONNX's boardcasting). If the scaler's shape is [1], no matter what shape the input is, we leave the axis
    # "None" because ONNX operator may automatically handle it. After ONNX-1.2, we adopt Numpy-style broadcasting rule.

    if target_opset < 7:
        # Input shape is [N, C, H, W]
        if len(shape) == 1:
            if shape[0] == 1:
                # shape is [] for scalar, we don't specify axis because it's a scalar
                return None, []
            else:
                # shape is [C], alignment starting at C-axis (indexed by 1)
                return 1, shape
        elif len(shape) == 3:
            if shape[0] == 1:
                # shape is [1, H, W], alignment starting at H-axis (indexed by 2)
                return 2, [shape[1], shape[2]]
            else:
                # shape is [C, H, W], alignment starting at C-axis (indexed by 1)
                return 1, shape
    else:
        # Input shape is [N, C, H, W]
        if len(shape) == 1:
            if shape[0] == 1:
                # shape is [1] for scalar, we don't specify axis because it's a scalar
                return None, [1]
            else:
                # shape is [C], alignment starting at C-axis
                return None, [shape[0], 1, 1]
        elif len(shape) == 3:
            if shape[0] == 1:
                # shape is [1, H, W], alignment starting at C-axis
                return None, shape
            else:
                # shape is [C, H, W], alignment starting at C-axis
                return None, shape


def convert_scale(scope, operator, container):
    # In CoreML's ScaleLayer, the input is first scaled by their "scale" attribute and then a "bias" can be added.
    # Symbols:
    #  a: scale attribute in CoreML's ScaleLayer
    #  b: bias attribute in CoreML's ScaleLayer
    #  x: input
    #  y: output
    # The math formulation of ScaleLayer should be
    #  y = a * x + b
    # Therefore, our strategy of composing ScaleLayer is to have one multiplication followed by an addition.
    params = operator.raw_operator.scale
    op1_type = 'Mul'
    scale_axis, scale_shape = deduce_broadcast_axis_and_shape(container.target_opset, params.shapeScale)
    scale_name = scope.get_unique_variable_name(op1_type + '_B')
    container.add_initializer(scale_name, onnx_proto.TensorProto.FLOAT, scale_shape, params.scale.floatValue)

    # CoreML is at most 3-D, so we always turn broadcasting on.
    scale_broadcast = 1

    if not params.hasBias:
        # Create a element-wise multiplication and use it to scale the input. The first input is the variable we want
        # to scale while the second input is their multipliers.
        apply_mul(scope, [operator.inputs[0].full_name, scale_name], operator.output_full_names, container,
                  operator_name=operator.full_name, axis=scale_axis, broadcast=scale_broadcast)
    else:
        # Declare a temporal variable to store the scaled input
        intra_variable_name = scope.get_unique_variable_name(operator.inputs[0].full_name + '_scaled')
        # Create a element-wise multiplication and use it to scale the input and save the result to a temporal variable
        apply_mul(scope, [operator.inputs[0].full_name, scale_name], intra_variable_name, container,
                  operator_name=operator.full_name, axis=scale_axis, broadcast=scale_broadcast)

        # Prepare materials to build an Add operator for adding bias
        bias_axis, bias_shape = deduce_broadcast_axis_and_shape(container.target_opset, params.shapeBias)

        # CoreML is at most 3-D, so we always turn broadcasting on.
        bias_broadcast = 1

        bias_name = scope.get_unique_variable_name(operator.full_name + '_B')
        container.add_initializer(bias_name, onnx_proto.TensorProto.FLOAT, bias_shape, params.bias.floatValue)
        # As bias exists, we add the bias into the output of the multiplication and then use the output of addition
        # as the final output of this conversion.
        apply_add(scope, [intra_variable_name, bias_name], operator.output_full_names, container,
                  axis=bias_axis, broadcast=bias_broadcast)


register_converter('scale', convert_scale)
