# SPDX-License-Identifier: Apache-2.0

from onnx import onnx_pb as onnx_proto
from ...common._apply_operation import apply_concat, apply_argmax
from ...common.data_types import Int64TensorType, FloatTensorType
from ...common.utils import check_input_and_output_numbers, check_input_and_output_types
from .linear_classifier import convert_sparkml_linear_classifier
from ..ops_names import get_sparkml_operator_name
from ...common._registration import register_converter, register_shape_calculator


def convert_one_vs_rest(scope, operator, container):
    op = operator.raw_operator
    classifier_output_names = []
    # initializer needed to extract the 2nd value of probability array
    index_tensor = scope.get_unique_variable_name("index_tensor")
    container.add_initializer(index_tensor, onnx_proto.TensorProto.INT64, [1], [1])
    # OneVsRest could have different classifiers
    # all must have at least Probability values for us to do the argmax
    for sub_model in op.models:
        classifier_op = scope.declare_local_operator(
            get_sparkml_operator_name(type(sub_model)), sub_model
        )
        classifier_op.raw_params = operator.raw_params
        classifier_op.inputs = operator.inputs
        classifier_prediction_output = scope.declare_local_variable(
            "classifier_prediction", FloatTensorType()
        )
        classifier_probability_output = scope.declare_local_variable(
            "classifier_probability", FloatTensorType()
        )
        classifier_op.outputs.append(classifier_prediction_output)
        classifier_op.outputs.append(classifier_probability_output)
        convert_sparkml_linear_classifier(scope, classifier_op, container)
        classifier_op.is_evaluated = True
        single_feature_tensor = scope.get_unique_variable_name("single_feature_tensor")
        container.add_node(
            "ArrayFeatureExtractor",
            [classifier_probability_output.full_name, index_tensor],
            single_feature_tensor,
            op_domain="ai.onnx.ml",
            name=scope.get_unique_operator_name("ArrayFeatureExtractor"),
        )
        classifier_output_names.append(single_feature_tensor)

    concatenated_probabilities = scope.get_unique_variable_name(
        "concatenated_predictions_tensor"
    )
    apply_concat(
        scope, classifier_output_names, concatenated_probabilities, container, axis=1
    )
    # to get Prediction from probability
    apply_argmax(
        scope,
        concatenated_probabilities,
        operator.outputs[0].full_name,
        container,
        axis=1,
        keepdims=1,
    )


register_converter("pyspark.ml.classification.OneVsRestModel", convert_one_vs_rest)


def calculate_one_vs_rest_output_shapes(operator):
    check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1)
    check_input_and_output_types(
        operator, good_input_types=[FloatTensorType, Int64TensorType]
    )
    if len(operator.inputs[0].type.shape) != 2:
        raise RuntimeError("Input must be a [N, C]-tensor")

    N = operator.inputs[0].type.shape[0]
    operator.outputs[0].type = Int64TensorType(shape=[N])


register_shape_calculator(
    "pyspark.ml.classification.OneVsRestModel", calculate_one_vs_rest_output_shapes
)
