# SPDX-License-Identifier: Apache-2.0

from .....proto import onnx_proto
from ....common._apply_operation import apply_sub
from ....common._registration import register_converter


def convert_preprocessing_mean_image(scope, operator, container):
    mean_tensor_name = scope.get_unique_variable_name(operator.full_name + "_mean")

    # We assume that the first input's shape is [N, C, H, W]
    # so that the mean image's shape, [C, H, W], can
    # be inferred from the first input's shape.
    container.add_initializer(
        mean_tensor_name,
        onnx_proto.TensorProto.FLOAT,
        operator.inputs[0].type.shape[1:],
        operator.raw_operator.meanImage,
    )

    # We assume that the first input variable's shape is [N, C, H, W]
    # while the mean image's shape is [C, H, W]. Thus,
    # broadcasting should be enabled starting with axis=1.
    apply_sub(
        scope,
        [operator.inputs[0].full_name, mean_tensor_name],
        operator.output_full_names,
        container,
        axis=1,
        broadcast=1,
    )


register_converter("meanImagePreprocessor", convert_preprocessing_mean_image)
