# SPDX-License-Identifier: Apache-2.0

import collections

from onnx import onnx_pb as onnx_proto
from ...common._apply_operation import apply_matmul, apply_exp, apply_add
from ...common._registration import register_converter, register_shape_calculator
from ...common.data_types import FloatTensorType
from ...common.utils import check_input_and_output_numbers


def convert_aft_survival_regression(scope, operator, container):
    op = operator.raw_operator

    coefficients = op.coefficients.toArray().astype(float)
    coefficients_tensor = scope.get_unique_variable_name("coefficients_tensor")
    container.add_initializer(
        coefficients_tensor,
        onnx_proto.TensorProto.FLOAT,
        [1, len(coefficients)],
        coefficients,
    )
    intercepts = (
        op.intercept.astype(float)
        if isinstance(op.intercept, collections.abc.Iterable)
        else [float(op.intercept)]
    )
    intercepts_tensor = scope.get_unique_variable_name("intercepts_tensor")
    container.add_initializer(
        intercepts_tensor, onnx_proto.TensorProto.FLOAT, [len(intercepts)], intercepts
    )

    matmul_result = scope.get_unique_variable_name("matmul_result_tensor")
    apply_matmul(
        scope,
        [operator.input_full_names[0], coefficients_tensor],
        matmul_result,
        container,
    )
    add_result = scope.get_unique_variable_name("intercept_added_tensor")
    apply_add(scope, [matmul_result, intercepts_tensor], add_result, container)
    apply_exp(scope, add_result, operator.output_full_names, container)


register_converter(
    "pyspark.ml.regression.AFTSurvivalRegressionModel", convert_aft_survival_regression
)


def calculate_aft_survival_regression_output_shapes(operator):
    check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1)

    N = operator.inputs[0].type.shape[0]
    operator.outputs[0].type = FloatTensorType([N, 1])


register_shape_calculator(
    "pyspark.ml.regression.AFTSurvivalRegressionModel",
    calculate_aft_survival_regression_output_shapes,
)
