# SPDX-License-Identifier: Apache-2.0

import collections

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_sparkml_linear_regressor(scope, operator, container):
    op = operator.raw_operator
    op_type = "LinearRegressor"
    attrs = {
        "name": scope.get_unique_operator_name(op_type),
        "coefficients": op.coefficients.astype(float),
        "intercepts": (
            op.intercept.astype(float)
            if isinstance(op.intercept, collections.abc.Iterable)
            else [float(op.intercept)]
        ),
    }
    container.add_node(
        op_type,
        operator.input_full_names,
        operator.output_full_names,
        op_domain="ai.onnx.ml",
        **attrs
    )


register_converter(
    "pyspark.ml.regression.LinearRegressionModel", convert_sparkml_linear_regressor
)
register_converter(
    "pyspark.ml.regression.GeneralizedLinearRegressionModel",
    convert_sparkml_linear_regressor,
)


def calculate_linear_regressor_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.LinearRegressionModel",
    calculate_linear_regressor_output_shapes,
)
register_shape_calculator(
    "pyspark.ml.regression.GeneralizedLinearRegressionModel",
    calculate_linear_regressor_output_shapes,
)
