# SPDX-License-Identifier: Apache-2.0

from .....proto import onnx_proto
from ....common._registration import register_converter
from ....common._apply_operation import apply_add
from .Scale import deduce_broadcast_axis_and_shape


def convert_bias(scope, operator, container):
    # Feed the input (which we are going to add a bias onto)
    # into Add operator. Its shape is [C, H, W] in CoreML but
    # [N, C, H, W] in ONNX.
    params = operator.raw_operator.bias

    # Adjust CoreML's bias shape and find a proper axis for broadcasting
    axis, shape = deduce_broadcast_axis_and_shape(container.target_opset, params.shape)

    # No matter what shape it is, we need "broadcast" on
    # because input shape is 4-D while bias is at most 3-D.
    broadcast = 1  # True

    # Create bias vector as an ONNX tensor
    bias_tensor_name = scope.get_unique_variable_name(operator.full_name + "_B")
    container.add_initializer(
        bias_tensor_name, onnx_proto.TensorProto.FLOAT, shape, params.bias.floatValue
    )

    apply_add(
        scope,
        [operator.inputs[0].full_name, bias_tensor_name],
        operator.output_full_names,
        container,
        operator_name=operator.full_name,
        axis=axis,
        broadcast=broadcast,
    )


register_converter("bias", convert_bias)
