# SPDX-License-Identifier: Apache-2.0

from .....proto import onnx_proto
from ....common._apply_operation import apply_batch_norm, apply_instance_norm
from ....common._registration import register_converter


def convert_batch_normalization(scope, operator, container):
    params = operator.raw_operator.batchnorm

    if params.instanceNormalization and not params.computeMeanVar:
        raise ValueError('It is impossible to do instance normalization without re-computing mean and variance')

    if params.instanceNormalization and params.computeMeanVar:
        op_type = 'InstanceNormalization'
    else:
        op_type = 'BatchNormalization'

    inputs = [operator.inputs[0].full_name]
    outputs = [operator.outputs[0].full_name]
    scale_tensor_name = scope.get_unique_variable_name(op_type + '_scale')
    container.add_initializer(scale_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels],
                              params.gamma.floatValue)
    inputs.append(scale_tensor_name)
    bias_tensor_name = scope.get_unique_variable_name(op_type + '_B')
    container.add_initializer(bias_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels], params.beta.floatValue)
    inputs.append(bias_tensor_name)

    epsilon = params.epsilon

    if op_type == 'BatchNormalization':
        mean_tensor_name = scope.get_unique_variable_name(op_type + '_mean')
        container.add_initializer(mean_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels],
                                  params.mean.floatValue)
        inputs.append(mean_tensor_name)
        variance_tensor_name = scope.get_unique_variable_name(op_type + '_variance')
        container.add_initializer(variance_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels],
                                  params.variance.floatValue)
        inputs.append(variance_tensor_name)
        momentum = 0.
        spatial = 1  # True

        if not params.instanceNormalization and params.computeMeanVar:
            # In this case, we apply batch normalization and adjust the statistics stored according the the batch
            # being processed.

            # To update "mean" and "var," we put their updated results back to the associated input tensors.
            outputs += inputs[1:3]
            # We also allocate two extra output buffers to store some intermediate results, but they are not used
            # in CoreML model.
            outputs.append(scope.get_unique_variable_name('saved_mean'))
            outputs.append(scope.get_unique_variable_name('saved_var'))
            # We choose "training" mode because some variables need to be updated.
            is_test = 0  # False
        elif not params.instanceNormalization and not params.computeMeanVar:
            # In this case, batch normalization is applied without updating mean, variance, etc. according to
            # the batches being processed. It means this operator works under testing model. Because there is no
            # variable update, we don't need to specify extra inputs and outputs like in previous code block.
            is_test = 1  # True
        else:
            raise ValueError('Unsupported operation mode')

        apply_batch_norm(scope, inputs, outputs, container, operator_name=operator.full_name, epsilon=epsilon,
                         is_test=is_test, momentum=momentum, spatial=spatial)
    else:
        apply_instance_norm(scope, inputs, outputs, container, operator_name=operator.full_name, epsilon=epsilon)


register_converter('batchnorm', convert_batch_normalization)
