# SPDX-License-Identifier: Apache-2.0

from ...common.data_types import StringTensorType
from ...common.utils import check_input_and_output_numbers, check_input_and_output_types
from ...common._registration import register_converter, register_shape_calculator


def convert_sparkml_tokenizer(scope, operator, container):
    op = operator.raw_operator
    if op is None:
        raise RuntimeError("op cannot be None.")
    # the SPARK version converts text to lowercase and applies "\\s" regexp to it
    # Here we'll tokenize and then normalize (to convert to lowercase)
    lowercase_output = scope.get_unique_variable_name('lowercase_tensor')
    container.add_node('StringNormalizer', operator.input_full_names[0], lowercase_output,
                       op_domain='ai.onnx',
                       name=scope.get_unique_operator_name('StringNormalizer'),
                       op_version=10,
                       case_change_action='LOWER')
    container.add_node('Tokenizer', lowercase_output, operator.output_full_names,
                       op_domain='com.microsoft',
                       name=scope.get_unique_operator_name('Tokenizer'),
                       op_version=1,
                       mark=0,
                       separators=[' ', '\t', '\r', '\n'],
                       pad_value='##ERROR##',
                       mincharnum=1)


register_converter('pyspark.ml.feature.Tokenizer', convert_sparkml_tokenizer)


def calculate_sparkml_tokenizer_output_shapes(operator):
    check_input_and_output_numbers(operator, output_count_range=1)
    check_input_and_output_types(operator,
                                 good_input_types=[StringTensorType])
    operator.outputs[0].type = StringTensorType()


register_shape_calculator('pyspark.ml.feature.Tokenizer', calculate_sparkml_tokenizer_output_shapes)
