# --------------------------------------------------------------------------
# ⚠️ WARNING - AUTO-GENERATED CODE - DO NOT EDIT ⚠️
# ⚙️ Generated by 'python -m opgen'
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
# pylint: disable=W0221,W0222,R0901,W0237
# mypy: disable-error-code=override
# ruff: noqa: N801
# --------------------------------------------------------------------------

from __future__ import annotations

from typing import Optional, Sequence, TypeVar, Union

from onnx.defs import get_schema
from typing_extensions import TypeAlias

from onnxscript.onnx_opset._impl.opset_ai_onnx_ml1 import Opset_ai_onnx_ml1
from onnxscript.onnx_types import FLOAT, INT64, STRING
from onnxscript.values import Op, Opset


class Opset_ai_onnx_ml2(Opset_ai_onnx_ml1):
    def __new__(cls):
        return Opset.__new__(cls, "ai.onnx.ml", 2)

    T1_LabelEncoder = TypeVar("T1_LabelEncoder", FLOAT, INT64, STRING)

    T2_LabelEncoder: TypeAlias = Union[FLOAT, INT64, STRING]

    def LabelEncoder(
        self,
        X: T1_LabelEncoder,
        *,
        default_float: float = -0.0,
        default_int64: int = -1,
        default_string: str = "_Unused",
        keys_floats: Optional[Sequence[float]] = None,
        keys_int64s: Optional[Sequence[int]] = None,
        keys_strings: Optional[Sequence[str]] = None,
        values_floats: Optional[Sequence[float]] = None,
        values_int64s: Optional[Sequence[int]] = None,
        values_strings: Optional[Sequence[str]] = None,
    ) -> T2_LabelEncoder:
        r"""[🌐 ai.onnx.ml::LabelEncoder(2)](https://onnx.ai/onnx/operators/onnx_aionnxml_LabelEncoder.html#labelencoder-2 "Online Documentation")


            Maps each element in the input tensor to another value.

            The mapping is determined by the two parallel attributes, 'keys_*' and
            'values_*' attribute. The i-th value in the specified 'keys_*' attribute
            would be mapped to the i-th value in the specified 'values_*' attribute. It
            implies that input's element type and the element type of the specified
            'keys_*' should be identical while the output type is identical to the
            specified 'values_*' attribute. If an input element can not be found in the
            specified 'keys_*' attribute, the 'default_*' that matches the specified
            'values_*' attribute may be used as its output value.

            Let's consider an example which maps a string tensor to an integer tensor.
            Assume and 'keys_strings' is ["Amy", "Sally"], 'values_int64s' is [5, 6],
            and 'default_int64' is '-1'.  The input ["Dori", "Amy", "Amy", "Sally",
            "Sally"] would be mapped to [-1, 5, 5, 6, 6].

            Since this operator is an one-to-one mapping, its input and output shapes
            are the same. Notice that only one of 'keys_*'/'values_*' can be set.

            For key look-up, bit-wise comparison is used so even a float NaN can be
            mapped to a value in 'values_*' attribute.



        Args:
            X: Input data. It can be either tensor or scalar.

            default_float: A float.

            default_int64: An integer.

            default_string: A string.

            keys_floats: A list of floats.

            keys_int64s: A list of ints.

            keys_strings: A list of strings. One and only one of 'keys_*'s should be
                set.

            values_floats: A list of floats.

            values_int64s: A list of ints.

            values_strings: A list of strings. One and only one of 'value_*'s should be
                set.
        """

        schema = get_schema("LabelEncoder", 2, "ai.onnx.ml")
        op = Op(self, "LabelEncoder", schema)
        return op(
            *self._prepare_inputs(schema, X),
            default_float=default_float,
            default_int64=default_int64,
            default_string=default_string,
            keys_floats=keys_floats,
            keys_int64s=keys_int64s,
            keys_strings=keys_strings,
            values_floats=values_floats,
            values_int64s=values_int64s,
            values_strings=values_strings,
        )