# SPDX-License-Identifier: Apache-2.0

from ...common._registration import register_shape_calculator
from ...common.data_types import DictionaryType, SequenceType, FloatTensorType
from ...common.utils import check_input_and_output_numbers, check_input_and_output_types


def calculate_dictionary_vectorizer_output_shapes(operator):
    """
    Allowed input/output patterns are
        1. Map ---> [1, C]

    C is the number of all allowed keys in the input dictionary.
    """
    # We assume all dictionaries' value types are float. It seems be reasonable to CoreML's
    # model input, but the existence of other map types leads to some concerns.
    check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1)
    # Two types are allowed. One is DictionaryType and
    # the other one, SequenceType, means a sequence of dictionaries.
    check_input_and_output_types(
        operator, good_input_types=[DictionaryType, SequenceType]
    )

    params = operator.raw_operator.dictVectorizer
    string_key_vector = params.stringToIndex.vector
    int64_key_vector = params.int64ToIndex.vector

    if len(string_key_vector) > 0 and len(int64_key_vector) > 0:
        raise RuntimeError("Only one key type can present at the same time")

    doc_string = operator.outputs[0].type.doc_string
    if len(string_key_vector) > 0:
        operator.outputs[0].type = FloatTensorType(
            [1, len(string_key_vector)], doc_string=doc_string
        )
    elif len(int64_key_vector) > 0:
        operator.outputs[1].type.shape = FloatTensorType(
            [1, len(int64_key_vector)], doc_string=doc_string
        )
    else:
        raise ValueError("Key vector cannot be empty")


register_shape_calculator(
    "dictVectorizer", calculate_dictionary_vectorizer_output_shapes
)
