# SPDX-FileCopyrightText: Copyright 2020-2022 Arm Limited and/or its affiliates <open-source-office@arm.com>
#
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the License); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an AS IS BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Description:
# TensorFlow Lite mapping functions used by both reader and writer.
# Contains a mapping from the various TensorFlow Lite enums and options structs, generated by the FlatBuffer code
# generator, to Vela's internal format.
import struct

import numpy as np

from .data_type import DataType
from .operation import CustomType
from .operation import Op
from .operation import Padding as opPad
from .operation import TensorIndices
from .tflite import AbsOptions
from .tflite import AddNOptions
from .tflite import AddOptions
from .tflite import ArgMaxOptions
from .tflite import ArgMinOptions
from .tflite import AssignVariableOptions
from .tflite import ATan2Options
from .tflite import BatchMatMulOptions
from .tflite import BatchToSpaceNDOptions
from .tflite import BidirectionalSequenceLSTMOptions
from .tflite import BidirectionalSequenceRNNOptions
from .tflite import BroadcastToOptions
from .tflite import BucketizeOptions
from .tflite import CallOnceOptions
from .tflite import CallOptions
from .tflite import CastOptions
from .tflite import ConcatEmbeddingsOptions
from .tflite import ConcatenationOptions
from .tflite import Conv2DOptions
from .tflite import Conv3DOptions
from .tflite import CosOptions
from .tflite import CumsumOptions
from .tflite import DensifyOptions
from .tflite import DepthToSpaceOptions
from .tflite import DepthwiseConv2DOptions
from .tflite import DequantizeOptions
from .tflite import DivOptions
from .tflite import DynamicUpdateSliceOptions
from .tflite import EmbeddingLookupSparseOptions
from .tflite import EqualOptions
from .tflite import ExpandDimsOptions
from .tflite import ExpOptions
from .tflite import FakeQuantOptions
from .tflite import FillOptions
from .tflite import FloorDivOptions
from .tflite import FloorModOptions
from .tflite import FullyConnectedOptions
from .tflite import GatherNdOptions
from .tflite import GatherOptions
from .tflite import GeluOptions
from .tflite import GreaterEqualOptions
from .tflite import GreaterOptions
from .tflite import HardSwishOptions
from .tflite import HashtableFindOptions
from .tflite import HashtableImportOptions
from .tflite import HashtableOptions
from .tflite import HashtableSizeOptions
from .tflite import IfOptions
from .tflite import L2NormOptions
from .tflite import LeakyReluOptions
from .tflite import LessEqualOptions
from .tflite import LessOptions
from .tflite import LocalResponseNormalizationOptions
from .tflite import LogicalAndOptions
from .tflite import LogicalNotOptions
from .tflite import LogicalOrOptions
from .tflite import LogSoftmaxOptions
from .tflite import LSHProjectionOptions
from .tflite import LSTMOptions
from .tflite import MatrixDiagOptions
from .tflite import MatrixSetDiagOptions
from .tflite import MaximumMinimumOptions
from .tflite import MirrorPadOptions
from .tflite import MulOptions
from .tflite import NegOptions
from .tflite import NonMaxSuppressionV4Options
from .tflite import NonMaxSuppressionV5Options
from .tflite import NotEqualOptions
from .tflite import OneHotOptions
from .tflite import PackOptions
from .tflite import PadOptions
from .tflite import PadV2Options
from .tflite import Pool2DOptions
from .tflite import PowOptions
from .tflite import QuantizeOptions
from .tflite import RandomOptions
from .tflite import RangeOptions
from .tflite import RankOptions
from .tflite import ReadVariableOptions
from .tflite import ReducerOptions
from .tflite import ReshapeOptions
from .tflite import ResizeBilinearOptions
from .tflite import ResizeNearestNeighborOptions
from .tflite import ReverseSequenceOptions
from .tflite import ReverseV2Options
from .tflite import Rfft2dOptions
from .tflite import RNNOptions
from .tflite import ScatterNdOptions
from .tflite import SegmentSumOptions
from .tflite import SelectOptions
from .tflite import SelectV2Options
from .tflite import SequenceRNNOptions
from .tflite import ShapeOptions
from .tflite import SignOptions
from .tflite import SkipGramOptions
from .tflite import SliceOptions
from .tflite import SoftmaxOptions
from .tflite import SpaceToBatchNDOptions
from .tflite import SpaceToDepthOptions
from .tflite import SparseToDenseOptions
from .tflite import SplitOptions
from .tflite import SplitVOptions
from .tflite import SquaredDifferenceOptions
from .tflite import SquareOptions
from .tflite import SqueezeOptions
from .tflite import StridedSliceOptions
from .tflite import SubOptions
from .tflite import SVDFOptions
from .tflite import TileOptions
from .tflite import TopKV2Options
from .tflite import TransposeConvOptions
from .tflite import TransposeOptions
from .tflite import UnidirectionalSequenceLSTMOptions
from .tflite import UniqueOptions
from .tflite import UnpackOptions
from .tflite import UnsortedSegmentMaxOptions
from .tflite import UnsortedSegmentMinOptions
from .tflite import UnsortedSegmentProdOptions
from .tflite import UnsortedSegmentSumOptions
from .tflite import VarHandleOptions
from .tflite import WhereOptions
from .tflite import WhileOptions
from .tflite import ZerosLikeOptions
from .tflite.ActivationFunctionType import ActivationFunctionType
from .tflite.BuiltinOperator import BuiltinOperator
from .tflite.BuiltinOptions import BuiltinOptions
from .tflite.Padding import Padding
from .tflite.TensorType import TensorType


def inverse_map(map):
    return {v: k for k, v in map.items()}


datatype_map = {
    TensorType.UINT8: DataType.uint8,
    TensorType.INT8: DataType.int8,
    TensorType.INT16: DataType.int16,
    TensorType.INT32: DataType.int32,
    TensorType.INT64: DataType.int64,
    TensorType.FLOAT16: DataType.float16,
    TensorType.FLOAT32: DataType.float32,
    TensorType.FLOAT64: DataType.float64,
    TensorType.STRING: DataType.string,
    TensorType.BOOL: DataType.bool,
    TensorType.COMPLEX64: DataType.complex64,
    TensorType.COMPLEX128: DataType.complex128,
    TensorType.UINT64: DataType.uint64,
    TensorType.RESOURCE: DataType.resource,
    TensorType.VARIANT: DataType.variant,
    TensorType.UINT32: DataType.uint32,
    TensorType.UINT16: DataType.uint16,
}

datatype_inv_map = inverse_map(datatype_map)
datatype_inv_map[DataType.quint8] = TensorType.UINT8

datatype_inv_map[DataType.qint8] = TensorType.INT8
datatype_inv_map[DataType.qint16] = TensorType.INT16
datatype_inv_map[DataType.qint32] = TensorType.INT32


datatype_map_numpy = {
    TensorType.UINT8: np.uint8,
    TensorType.UINT16: np.uint16,
    TensorType.UINT32: np.uint32,
    TensorType.UINT64: np.uint64,
    TensorType.INT8: np.int8,
    TensorType.INT16: np.int16,
    TensorType.INT32: np.int32,
    TensorType.INT64: np.int64,
    TensorType.FLOAT16: np.float16,
    TensorType.FLOAT32: np.float32,
    TensorType.FLOAT64: np.float64,
    TensorType.BOOL: np.bool,
    TensorType.COMPLEX64: np.complex64,
    TensorType.COMPLEX128: np.complex128,
    TensorType.STRING: np.uint8,
    TensorType.RESOURCE: None,  # experimental type
    TensorType.VARIANT: None,  # experimental type
}


builtin_options_map = {
    BuiltinOptions.Conv2DOptions: Conv2DOptions.Conv2DOptions,
    BuiltinOptions.DepthwiseConv2DOptions: DepthwiseConv2DOptions.DepthwiseConv2DOptions,
    BuiltinOptions.ConcatEmbeddingsOptions: ConcatEmbeddingsOptions.ConcatEmbeddingsOptions,
    BuiltinOptions.LSHProjectionOptions: LSHProjectionOptions.LSHProjectionOptions,
    BuiltinOptions.Pool2DOptions: Pool2DOptions.Pool2DOptions,
    BuiltinOptions.SVDFOptions: SVDFOptions.SVDFOptions,
    BuiltinOptions.RNNOptions: RNNOptions.RNNOptions,
    BuiltinOptions.FullyConnectedOptions: FullyConnectedOptions.FullyConnectedOptions,
    BuiltinOptions.SoftmaxOptions: SoftmaxOptions.SoftmaxOptions,
    BuiltinOptions.ConcatenationOptions: ConcatenationOptions.ConcatenationOptions,
    BuiltinOptions.AddOptions: AddOptions.AddOptions,
    BuiltinOptions.L2NormOptions: L2NormOptions.L2NormOptions,
    BuiltinOptions.LocalResponseNormalizationOptions: LocalResponseNormalizationOptions.LocalResponseNormalizationOptions,  # noqa: E501
    BuiltinOptions.LSTMOptions: LSTMOptions.LSTMOptions,
    BuiltinOptions.ResizeBilinearOptions: ResizeBilinearOptions.ResizeBilinearOptions,
    BuiltinOptions.CallOptions: CallOptions.CallOptions,
    BuiltinOptions.ReshapeOptions: ReshapeOptions.ReshapeOptions,
    BuiltinOptions.SkipGramOptions: SkipGramOptions.SkipGramOptions,
    BuiltinOptions.SpaceToDepthOptions: SpaceToDepthOptions.SpaceToDepthOptions,
    BuiltinOptions.EmbeddingLookupSparseOptions: EmbeddingLookupSparseOptions.EmbeddingLookupSparseOptions,
    BuiltinOptions.MulOptions: MulOptions.MulOptions,
    BuiltinOptions.PadOptions: PadOptions.PadOptions,
    BuiltinOptions.GatherOptions: GatherOptions.GatherOptions,
    BuiltinOptions.BatchToSpaceNDOptions: BatchToSpaceNDOptions.BatchToSpaceNDOptions,
    BuiltinOptions.SpaceToBatchNDOptions: SpaceToBatchNDOptions.SpaceToBatchNDOptions,
    BuiltinOptions.TransposeOptions: TransposeOptions.TransposeOptions,
    BuiltinOptions.ReducerOptions: ReducerOptions.ReducerOptions,
    BuiltinOptions.SubOptions: SubOptions.SubOptions,
    BuiltinOptions.DivOptions: DivOptions.DivOptions,
    BuiltinOptions.SqueezeOptions: SqueezeOptions.SqueezeOptions,
    BuiltinOptions.SequenceRNNOptions: SequenceRNNOptions.SequenceRNNOptions,
    BuiltinOptions.StridedSliceOptions: StridedSliceOptions.StridedSliceOptions,
    BuiltinOptions.ExpOptions: ExpOptions.ExpOptions,
    BuiltinOptions.TopKV2Options: TopKV2Options.TopKV2Options,
    BuiltinOptions.SplitOptions: SplitOptions.SplitOptions,
    BuiltinOptions.LogSoftmaxOptions: LogSoftmaxOptions.LogSoftmaxOptions,
    BuiltinOptions.CastOptions: CastOptions.CastOptions,
    BuiltinOptions.DequantizeOptions: DequantizeOptions.DequantizeOptions,
    BuiltinOptions.MaximumMinimumOptions: MaximumMinimumOptions.MaximumMinimumOptions,
    BuiltinOptions.ArgMaxOptions: ArgMaxOptions.ArgMaxOptions,
    BuiltinOptions.LessOptions: LessOptions.LessOptions,
    BuiltinOptions.NegOptions: NegOptions.NegOptions,
    BuiltinOptions.PadV2Options: PadV2Options.PadV2Options,
    BuiltinOptions.GreaterOptions: GreaterOptions.GreaterOptions,
    BuiltinOptions.GreaterEqualOptions: GreaterEqualOptions.GreaterEqualOptions,
    BuiltinOptions.LessEqualOptions: LessEqualOptions.LessEqualOptions,
    BuiltinOptions.SelectOptions: SelectOptions.SelectOptions,
    BuiltinOptions.SliceOptions: SliceOptions.SliceOptions,
    BuiltinOptions.TransposeConvOptions: TransposeConvOptions.TransposeConvOptions,
    BuiltinOptions.SparseToDenseOptions: SparseToDenseOptions.SparseToDenseOptions,
    BuiltinOptions.TileOptions: TileOptions.TileOptions,
    BuiltinOptions.ExpandDimsOptions: ExpandDimsOptions.ExpandDimsOptions,
    BuiltinOptions.EqualOptions: EqualOptions.EqualOptions,
    BuiltinOptions.NotEqualOptions: NotEqualOptions.NotEqualOptions,
    BuiltinOptions.ShapeOptions: ShapeOptions.ShapeOptions,
    BuiltinOptions.PowOptions: PowOptions.PowOptions,
    BuiltinOptions.ArgMinOptions: ArgMinOptions.ArgMinOptions,
    BuiltinOptions.FakeQuantOptions: FakeQuantOptions.FakeQuantOptions,
    BuiltinOptions.PackOptions: PackOptions.PackOptions,
    BuiltinOptions.LogicalOrOptions: LogicalOrOptions.LogicalOrOptions,
    BuiltinOptions.OneHotOptions: OneHotOptions.OneHotOptions,
    BuiltinOptions.LogicalAndOptions: LogicalAndOptions.LogicalAndOptions,
    BuiltinOptions.LogicalNotOptions: LogicalNotOptions.LogicalNotOptions,
    BuiltinOptions.UnpackOptions: UnpackOptions.UnpackOptions,
    BuiltinOptions.FloorDivOptions: FloorDivOptions.FloorDivOptions,
    BuiltinOptions.SquareOptions: SquareOptions.SquareOptions,
    BuiltinOptions.ZerosLikeOptions: ZerosLikeOptions.ZerosLikeOptions,
    BuiltinOptions.FillOptions: FillOptions.FillOptions,
    BuiltinOptions.BidirectionalSequenceLSTMOptions: BidirectionalSequenceLSTMOptions.BidirectionalSequenceLSTMOptions,
    BuiltinOptions.BidirectionalSequenceRNNOptions: BidirectionalSequenceRNNOptions.BidirectionalSequenceRNNOptions,
    BuiltinOptions.UnidirectionalSequenceLSTMOptions: UnidirectionalSequenceLSTMOptions.UnidirectionalSequenceLSTMOptions,  # noqa: E501
    BuiltinOptions.FloorModOptions: FloorModOptions.FloorModOptions,
    BuiltinOptions.RangeOptions: RangeOptions.RangeOptions,
    BuiltinOptions.ResizeNearestNeighborOptions: ResizeNearestNeighborOptions.ResizeNearestNeighborOptions,
    BuiltinOptions.LeakyReluOptions: LeakyReluOptions.LeakyReluOptions,
    BuiltinOptions.SquaredDifferenceOptions: SquaredDifferenceOptions.SquaredDifferenceOptions,
    BuiltinOptions.MirrorPadOptions: MirrorPadOptions.MirrorPadOptions,
    BuiltinOptions.AbsOptions: AbsOptions.AbsOptions,
    BuiltinOptions.SplitVOptions: SplitVOptions.SplitVOptions,
    BuiltinOptions.UniqueOptions: UniqueOptions.UniqueOptions,
    BuiltinOptions.ReverseV2Options: ReverseV2Options.ReverseV2Options,
    BuiltinOptions.AddNOptions: AddNOptions.AddNOptions,
    BuiltinOptions.GatherNdOptions: GatherNdOptions.GatherNdOptions,
    BuiltinOptions.CosOptions: CosOptions.CosOptions,
    BuiltinOptions.WhereOptions: WhereOptions.WhereOptions,
    BuiltinOptions.RankOptions: RankOptions.RankOptions,
    BuiltinOptions.ReverseSequenceOptions: ReverseSequenceOptions.ReverseSequenceOptions,
    BuiltinOptions.MatrixDiagOptions: MatrixDiagOptions.MatrixDiagOptions,
    BuiltinOptions.QuantizeOptions: QuantizeOptions.QuantizeOptions,
    BuiltinOptions.MatrixSetDiagOptions: MatrixSetDiagOptions.MatrixSetDiagOptions,
    BuiltinOptions.DensifyOptions: DensifyOptions.DensifyOptions,
    BuiltinOptions.DepthToSpaceOptions: DepthToSpaceOptions.DepthToSpaceOptions,
    BuiltinOptions.HardSwishOptions: HardSwishOptions.HardSwishOptions,
    BuiltinOptions.IfOptions: IfOptions.IfOptions,
    BuiltinOptions.NonMaxSuppressionV4Options: NonMaxSuppressionV4Options.NonMaxSuppressionV4Options,
    BuiltinOptions.NonMaxSuppressionV5Options: NonMaxSuppressionV5Options.NonMaxSuppressionV5Options,
    BuiltinOptions.ScatterNdOptions: ScatterNdOptions.ScatterNdOptions,
    BuiltinOptions.SegmentSumOptions: SegmentSumOptions.SegmentSumOptions,
    BuiltinOptions.SelectV2Options: SelectV2Options.SelectV2Options,
    BuiltinOptions.WhileOptions: WhileOptions.WhileOptions,
    BuiltinOptions.BatchMatMulOptions: BatchMatMulOptions.BatchMatMulOptions,
    BuiltinOptions.CumsumOptions: CumsumOptions.CumsumOptions,
    BuiltinOptions.CallOnceOptions: CallOnceOptions.CallOnceOptions,
    BuiltinOptions.BroadcastToOptions: BroadcastToOptions.BroadcastToOptions,
    BuiltinOptions.Rfft2dOptions: Rfft2dOptions.Rfft2dOptions,
    BuiltinOptions.Conv3DOptions: Conv3DOptions.Conv3DOptions,
    BuiltinOptions.HashtableOptions: HashtableOptions.HashtableOptions,
    BuiltinOptions.HashtableFindOptions: HashtableFindOptions.HashtableFindOptions,
    BuiltinOptions.HashtableImportOptions: HashtableImportOptions.HashtableImportOptions,
    BuiltinOptions.HashtableSizeOptions: HashtableSizeOptions.HashtableSizeOptions,
    BuiltinOptions.VarHandleOptions: VarHandleOptions.VarHandleOptions,
    BuiltinOptions.ReadVariableOptions: ReadVariableOptions.ReadVariableOptions,
    BuiltinOptions.AssignVariableOptions: AssignVariableOptions.AssignVariableOptions,
    BuiltinOptions.RandomOptions: RandomOptions.RandomOptions,
    BuiltinOptions.BucketizeOptions: BucketizeOptions.BucketizeOptions,
    BuiltinOptions.DynamicUpdateSliceOptions: DynamicUpdateSliceOptions.DynamicUpdateSliceOptions,
    BuiltinOptions.GeluOptions: GeluOptions.GeluOptions,
    BuiltinOptions.UnsortedSegmentProdOptions: UnsortedSegmentProdOptions.UnsortedSegmentProdOptions,
    BuiltinOptions.UnsortedSegmentMaxOptions: UnsortedSegmentMaxOptions.UnsortedSegmentMaxOptions,
    BuiltinOptions.UnsortedSegmentMinOptions: UnsortedSegmentMinOptions.UnsortedSegmentMinOptions,
    BuiltinOptions.UnsortedSegmentSumOptions: UnsortedSegmentSumOptions.UnsortedSegmentSumOptions,
    BuiltinOptions.ATan2Options: ATan2Options.ATan2Options,
    BuiltinOptions.SignOptions: SignOptions.SignOptions,
}


builtin_options_inv_map = inverse_map(builtin_options_map)


def underscore_to_camel_case(s):
    return "".join(x.title() for x in s.split("_"))


def padding_deserialize(x):
    return padding_map[x]


def padding_serialize(builder, x):
    return padding_inv_map[x]


def activation_deserialize(x):
    return activation_function_map[x]


def activation_serialize(builder, x):
    return activation_function_inv_map[x]


def datatype_deserialize(x):
    return datatype_map[x]


def datatype_serialize(builder, x):
    return datatype_inv_map[x]


def identity(x):
    return x


def identity_serialize(builder, x):
    return x


def write_byte_vector(builder, v):
    builder.StartVector(1, len(v), 1)
    for e in v[::-1]:
        builder.PrependByte(e)
    return builder.EndVector()


def write_int_vector(builder, v):
    builder.StartVector(4, len(v), 4)
    for e in v[::-1]:
        builder.PrependInt32(e)
    return builder.EndVector()


class OptionsSerializer:
    def __init__(self, name, members=None):
        self.name = name
        self.module = globals()[self.name]
        self.cls = getattr(self.module, self.name)
        self.builtin_opt_type = builtin_options_inv_map[self.cls]
        self.members = []
        if members is not None:
            for mem in members:
                deserialize = identity
                serialize = identity_serialize
                is_vector = False
                if isinstance(mem, tuple):
                    if len(mem) == 3:
                        mem, deserialize, serialize = mem
                    elif len(mem) == 2:
                        mem, is_vector = mem
                        deserialize = tuple
                        serialize = write_int_vector
                    else:
                        assert 0
                underscore_mem = mem
                camelcase_mem = underscore_to_camel_case(mem)
                self.members.append((underscore_mem, camelcase_mem, deserialize, serialize, is_vector))

    def deserialize(self, op_data):
        builtin_options = op_data.BuiltinOptions()
        attrs = {}
        if builtin_options:
            tfattrs = self.cls()
            tfattrs.Init(builtin_options.Bytes, builtin_options.Pos)
            for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
                fun = camelcase_mem
                if is_vector:
                    fun += "AsNumpy"

                attr = getattr(tfattrs, fun)()
                try:
                    attrs[underscore_mem] = deserialize(attr)
                except TypeError:
                    print("Warning: {0} could not read attribute '{1}'.".format(self.name, underscore_mem))

        return attrs

    def serialize(self, builder, attrs):
        ser_attrs = []
        for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
            a = serialize(builder, attrs[underscore_mem])
            ser_attrs.append((camelcase_mem, a))

        getattr(self.module, self.name + "Start")(builder)

        for camelcase_mem, a in ser_attrs:
            getattr(self.module, self.name + "Add" + camelcase_mem)(builder, a)

        return getattr(self.module, self.name + "End")(builder), None


class CustomOptionsSerializer:
    CUSTOM_OPTIONS_NPU_OP = [0x01, 0x04, 0x01]  # NpuOp=1, FlexbufferFormat.UINT8=4, byte length=1
    CUSTOM_OPTIONS_FORMAT_DEFAULT = 0

    def __init__(self):
        self.custom_opt_format = 0

    def deserialize(self, op_data):
        attrs = {}
        custom_options = op_data.CustomOptionsAsNumpy()
        attrs["custom_options"] = custom_options
        attrs["custom_options_format"] = op_data.CustomOptionsFormat()

        if np.array_equal(custom_options, self.CUSTOM_OPTIONS_NPU_OP):
            attrs["custom_type"] = CustomType.ExistingNpuOp

        return attrs

    def serialize(self, builder, attrs):
        custom_type = attrs.get("custom_type", CustomType.ThirdPartyOp)
        self.custom_opt_format = attrs.get("custom_options_format", self.CUSTOM_OPTIONS_FORMAT_DEFAULT)

        # Set NPU op custom options for the TensorFlow Lite custom operator
        if custom_type == CustomType.NpuOp:
            custom_options = self.CUSTOM_OPTIONS_NPU_OP
        else:
            custom_options = attrs.get("custom_options", [])

        custom_options_bytes = struct.pack("<{0}B".format(len(custom_options)), *custom_options)
        custom_offset = write_byte_vector(builder, custom_options_bytes)

        return None, custom_offset


padding_map = {
    Padding.SAME: opPad.SAME,
    Padding.VALID: opPad.VALID,
}

padding_inv_map = inverse_map(padding_map)


activation_function_map = {
    ActivationFunctionType.NONE: None,
    ActivationFunctionType.RELU: Op.Relu,
    ActivationFunctionType.RELU_N1_TO_1: Op.ReluN1To1,
    ActivationFunctionType.RELU6: Op.Relu6,
    ActivationFunctionType.TANH: Op.Tanh,
    ActivationFunctionType.SIGN_BIT: Op.SignBit,
}

activation_function_inv_map = inverse_map(activation_function_map)

fused_act = ("fused_activation_function", activation_deserialize, activation_serialize)
padding = ("padding", padding_deserialize, padding_serialize)

is_int_vec = True

TFLITE_NO_INDICES = TensorIndices([], [], [])
TFLITE_IFM_INDICES = TensorIndices([0], [], [])
TFLITE_IFM_WEIGHTS_INDICES = TensorIndices([0], [1], [])
TFLITE_IFM_WEIGHTS_BIAS_INDICES = TensorIndices([0], [1], [2])
TFLITE_IFM_IFM2_INDICES = TensorIndices([0, 1], [], [])
TFLITE_CONV2D_BACKPROP_INDICES = TensorIndices([2], [1], [3])
TFLITE_TRANSPOSE_CONV_INDICES = TensorIndices([0], [1], [3])
TFLITE_CONCAT_INDICES = TensorIndices([1, 2], [], [])
TFLITE_SPLIT_IFM_INDICES = TensorIndices([1], [], [])
TFLITE_BLOCK_LSTM_INDICES = TensorIndices([3], [4], [])

builtin_operator_map = {
    BuiltinOperator.ADD: (
        Op.Add,
        OptionsSerializer("AddOptions", (fused_act, "pot_scale_int16")),
        TFLITE_IFM_IFM2_INDICES,
    ),
    BuiltinOperator.AVERAGE_POOL_2D: (
        Op.AvgPool,
        OptionsSerializer(
            "Pool2DOptions", ("filter_height", "filter_width", fused_act, padding, "stride_h", "stride_w")
        ),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.ATAN2: (Op.Atan2, OptionsSerializer("ATan2Options", ()), TFLITE_IFM_IFM2_INDICES),
    BuiltinOperator.CONCATENATION: (
        Op.ConcatTFLite,
        OptionsSerializer("ConcatenationOptions", ("axis", fused_act)),
        TFLITE_CONCAT_INDICES,
    ),
    BuiltinOperator.CONV_2D: (
        Op.Conv2DBias,
        OptionsSerializer(
            "Conv2DOptions", ("dilation_h_factor", "dilation_w_factor", fused_act, padding, "stride_h", "stride_w")
        ),
        TFLITE_IFM_WEIGHTS_BIAS_INDICES,
    ),
    BuiltinOperator.DEPTHWISE_CONV_2D: (
        Op.DepthwiseConv2DBias,
        OptionsSerializer(
            "DepthwiseConv2DOptions",
            ("depth_multiplier", "dilation_h_factor", "dilation_w_factor", fused_act, padding, "stride_h", "stride_w"),
        ),
        TFLITE_IFM_WEIGHTS_BIAS_INDICES,
    ),
    BuiltinOperator.DEPTH_TO_SPACE: (
        Op.DepthToSpace,
        OptionsSerializer("DepthToSpaceOptions", ("block_size",)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.DEQUANTIZE: (Op.Dequantize, OptionsSerializer("DequantizeOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.EMBEDDING_LOOKUP: (Op.EmbeddingLookup, None, TFLITE_NO_INDICES),
    BuiltinOperator.FLOOR: (Op.Floor, None, TFLITE_NO_INDICES),
    BuiltinOperator.FULLY_CONNECTED: (
        Op.FullyConnected,
        OptionsSerializer(
            "FullyConnectedOptions", ("asymmetric_quantize_inputs", fused_act, "keep_num_dims", "weights_format")
        ),
        TFLITE_IFM_WEIGHTS_BIAS_INDICES,
    ),
    BuiltinOperator.HASHTABLE_LOOKUP: (Op.HashtableLookup, None, TFLITE_NO_INDICES),
    BuiltinOperator.L2_NORMALIZATION: (Op.L2Norm, OptionsSerializer("L2NormOptions", (fused_act,)), TFLITE_NO_INDICES),
    BuiltinOperator.L2_POOL_2D: (
        Op.L2Pool2D,
        OptionsSerializer(
            "Pool2DOptions", ("filter_height", "filter_width", fused_act, padding, "stride_h", "stride_w")
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.LOCAL_RESPONSE_NORMALIZATION: (
        Op.LRN,
        OptionsSerializer("LocalResponseNormalizationOptions", ("alpha", "beta", "bias", "radius")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.LOGISTIC: (Op.Sigmoid, None, TFLITE_IFM_INDICES),
    BuiltinOperator.LSH_PROJECTION: (
        Op.LSHProjection,
        OptionsSerializer("LSHProjectionOptions", ("type",)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.LSTM: (
        Op.Lstm,
        OptionsSerializer(
            "LSTMOptions", ("asymmetric_quantize_inputs", "cell_clip", fused_act, "kernel_type", "proj_clip")
        ),
        TFLITE_IFM_WEIGHTS_INDICES,
    ),
    BuiltinOperator.MAX_POOL_2D: (
        Op.MaxPool,
        OptionsSerializer(
            "Pool2DOptions", ("filter_height", "filter_width", fused_act, padding, "stride_h", "stride_w")
        ),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.MUL: (Op.Mul, OptionsSerializer("MulOptions", (fused_act,)), TFLITE_IFM_IFM2_INDICES),
    BuiltinOperator.RELU: (Op.Relu, None, TFLITE_IFM_INDICES),
    BuiltinOperator.RELU_0_TO_1: (Op.Relu0To1, None, TFLITE_IFM_INDICES),
    BuiltinOperator.RELU_N1_TO_1: (Op.ReluN1To1, None, TFLITE_IFM_INDICES),
    BuiltinOperator.RELU6: (Op.Relu6, None, TFLITE_IFM_INDICES),
    BuiltinOperator.RESHAPE: (
        Op.Reshape,
        OptionsSerializer("ReshapeOptions", (("new_shape", is_int_vec),)),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.RESIZE_BILINEAR: (
        Op.ResizeBilinear,
        OptionsSerializer("ResizeBilinearOptions", ("align_corners", "half_pixel_centers")),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.RNN: (
        Op.Rnn,
        OptionsSerializer("RNNOptions", ("asymmetric_quantize_inputs", fused_act)),
        TFLITE_IFM_WEIGHTS_INDICES,
    ),
    BuiltinOperator.SOFTMAX: (Op.Softmax, OptionsSerializer("SoftmaxOptions", ("beta",)), TFLITE_IFM_INDICES),
    BuiltinOperator.SPACE_TO_DEPTH: (
        Op.SpaceToDepth,
        OptionsSerializer("SpaceToDepthOptions", ("block_size",)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.SVDF: (
        Op.Svdf,
        OptionsSerializer("SVDFOptions", ("asymmetric_quantize_inputs", fused_act, "rank")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.TANH: (Op.Tanh, None, TFLITE_IFM_INDICES),
    BuiltinOperator.CONCAT_EMBEDDINGS: (
        Op.ConcatEmbeddings,
        OptionsSerializer(
            "ConcatEmbeddingsOptions",
            (
                "embedding_dim_per_channel",
                "embedding_dim_per_channel_as_numpy",
                "embedding_dim_per_channel_is_none",
                "embedding_dim_per_channel_length",
                "num_channels",
                "num_columns_per_channel",
                "num_columns_per_channel_as_numpy",
                "num_columns_per_channel_is_none",
                "num_columns_per_channel_length",
            ),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.SKIP_GRAM: (
        Op.SkipGram,
        OptionsSerializer("SkipGramOptions", ("include_all_ngrams", "max_skip_size", "ngram_size")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.CALL: (Op.Call, OptionsSerializer("CallOptions", ("subgraph",)), TFLITE_NO_INDICES),
    BuiltinOperator.EMBEDDING_LOOKUP_SPARSE: (
        Op.EmbeddingLookupSparse,
        OptionsSerializer("EmbeddingLookupSparseOptions", ("combiner",)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.PAD: (Op.Pad, OptionsSerializer("PadOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_RNN: (
        Op.UnidirectionalSequenceRnn,
        OptionsSerializer("SequenceRNNOptions", ("asymmetric_quantize_inputs", fused_act, "time_major")),
        TFLITE_IFM_WEIGHTS_INDICES,
    ),
    BuiltinOperator.GATHER: (
        Op.GatherV2,
        OptionsSerializer("GatherOptions", ("axis", "batch_dims")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.BATCH_TO_SPACE_ND: (
        Op.BatchToSpaceND,
        OptionsSerializer("BatchToSpaceNDOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.SPACE_TO_BATCH_ND: (
        Op.SpaceToBatchND,
        OptionsSerializer("SpaceToBatchNDOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.TRANSPOSE: (Op.Transpose, OptionsSerializer("TransposeOptions"), TFLITE_IFM_IFM2_INDICES),
    BuiltinOperator.MEAN: (Op.Mean, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_IFM_INDICES),
    BuiltinOperator.SUB: (
        Op.Sub,
        OptionsSerializer("SubOptions", (fused_act, "pot_scale_int16")),
        TFLITE_IFM_IFM2_INDICES,
    ),
    BuiltinOperator.DIV: (Op.Div, OptionsSerializer("DivOptions", (fused_act,)), TFLITE_NO_INDICES),
    BuiltinOperator.SQUEEZE: (
        Op.Squeeze,
        OptionsSerializer(
            "SqueezeOptions",
            (("squeeze_dims", is_int_vec),),
        ),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_LSTM: (
        Op.UnidirectionalSequenceLstm,
        OptionsSerializer(
            "UnidirectionalSequenceLSTMOptions",
            ("asymmetric_quantize_inputs", "cell_clip", fused_act, "proj_clip", "time_major"),
        ),
        TFLITE_IFM_WEIGHTS_INDICES,
    ),
    BuiltinOperator.STRIDED_SLICE: (
        Op.StridedSlice,
        OptionsSerializer(
            "StridedSliceOptions", ("begin_mask", "ellipsis_mask", "end_mask", "new_axis_mask", "shrink_axis_mask")
        ),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.BIDIRECTIONAL_SEQUENCE_RNN: (
        Op.BidirectionalSequenceRnn,
        OptionsSerializer(
            "BidirectionalSequenceRNNOptions", ("asymmetric_quantize_inputs", fused_act, "merge_outputs", "time_major")
        ),
        TFLITE_IFM_WEIGHTS_INDICES,
    ),
    BuiltinOperator.EXP: (Op.Exp, OptionsSerializer("ExpOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.TOPK_V2: (Op.TopKV2, OptionsSerializer("TopKV2Options"), TFLITE_NO_INDICES),
    BuiltinOperator.SPLIT: (Op.Split, OptionsSerializer("SplitOptions", ("num_splits",)), TFLITE_SPLIT_IFM_INDICES),
    BuiltinOperator.LOG_SOFTMAX: (Op.LogSoftmax, OptionsSerializer("LogSoftmaxOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.DELEGATE: (Op.Delegate, None, TFLITE_NO_INDICES),
    BuiltinOperator.BIDIRECTIONAL_SEQUENCE_LSTM: (
        Op.BidirectionalSequenceLstm,
        OptionsSerializer(
            "BidirectionalSequenceLSTMOptions",
            ("asymmetric_quantize_inputs", "cell_clip", fused_act, "merge_outputs", "proj_clip", "time_major"),
        ),
        TFLITE_IFM_WEIGHTS_INDICES,
    ),
    BuiltinOperator.CAST: (
        Op.Cast,
        OptionsSerializer(
            "CastOptions",
            (
                ("in_data_type", datatype_deserialize, datatype_serialize),
                ("out_data_type", datatype_deserialize, datatype_serialize),
            ),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.PRELU: (Op.Prelu, None, TFLITE_IFM_IFM2_INDICES),
    BuiltinOperator.MAXIMUM: (Op.Maximum, OptionsSerializer("MaximumMinimumOptions"), TFLITE_IFM_IFM2_INDICES),
    BuiltinOperator.ARG_MAX: (
        Op.ArgMax,
        OptionsSerializer("ArgMaxOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.MINIMUM: (Op.Minimum, OptionsSerializer("MaximumMinimumOptions"), TFLITE_IFM_IFM2_INDICES),
    BuiltinOperator.LESS: (Op.Less, OptionsSerializer("LessOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.NEG: (Op.Neg, OptionsSerializer("NegOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.PADV2: (Op.PadV2, OptionsSerializer("PadV2Options"), TFLITE_NO_INDICES),
    BuiltinOperator.GREATER: (Op.Greater, OptionsSerializer("GreaterOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.GREATER_EQUAL: (Op.GreaterEqual, OptionsSerializer("GreaterEqualOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.LESS_EQUAL: (Op.LessEqual, OptionsSerializer("LessEqualOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.SELECT: (Op.Select, OptionsSerializer("SelectOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.SLICE: (Op.Slice, OptionsSerializer("SliceOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.SIN: (Op.Sin, None, TFLITE_NO_INDICES),
    BuiltinOperator.TRANSPOSE_CONV: (
        Op.Conv2DBackpropInput,
        OptionsSerializer("TransposeConvOptions", (padding, "stride_h", "stride_w")),
        TFLITE_CONV2D_BACKPROP_INDICES,
    ),
    BuiltinOperator.SIGN: (Op.Sign, OptionsSerializer("SignOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.SPARSE_TO_DENSE: (
        Op.SparseToDense,
        OptionsSerializer("SparseToDenseOptions", ("validate_indices",)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.TILE: (Op.Tile, OptionsSerializer("TileOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.EXPAND_DIMS: (Op.ExpandDims, OptionsSerializer("ExpandDimsOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.EQUAL: (Op.Equal, OptionsSerializer("EqualOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.NOT_EQUAL: (Op.NotEqual, OptionsSerializer("NotEqualOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.LOG: (Op.Log, None, TFLITE_NO_INDICES),
    BuiltinOperator.SUM: (Op.Sum, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
    BuiltinOperator.SQRT: (Op.Sqrt, None, TFLITE_NO_INDICES),
    BuiltinOperator.RSQRT: (Op.Rsqrt, None, TFLITE_NO_INDICES),
    BuiltinOperator.SHAPE: (
        Op.Shape,
        OptionsSerializer("ShapeOptions", (("out_type", datatype_deserialize, datatype_serialize),)),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.POW: (Op.Pow, OptionsSerializer("PowOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.ARG_MIN: (
        Op.ArgMin,
        OptionsSerializer("ArgMinOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.FAKE_QUANT: (
        Op.FakeQuantWithMinMaxArgs,
        OptionsSerializer("FakeQuantOptions", ("max", "min", "narrow_range", "num_bits")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.REDUCE_PROD: (Op.Prod, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
    BuiltinOperator.REDUCE_MAX: (Op.Max, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
    BuiltinOperator.PACK: (Op.Pack, OptionsSerializer("PackOptions", ("axis", "values_count")), TFLITE_IFM_INDICES),
    BuiltinOperator.LOGICAL_OR: (Op.LogicalOr, OptionsSerializer("LogicalOrOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.ONE_HOT: (Op.OneHot, OptionsSerializer("OneHotOptions", ("axis",)), TFLITE_NO_INDICES),
    BuiltinOperator.LOGICAL_AND: (Op.LogicalAnd, OptionsSerializer("LogicalAndOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.LOGICAL_NOT: (Op.LogicalNot, OptionsSerializer("LogicalNotOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.UNPACK: (Op.Unpack, OptionsSerializer("UnpackOptions", ("axis", "num")), TFLITE_IFM_INDICES),
    BuiltinOperator.REDUCE_MIN: (Op.Min, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
    BuiltinOperator.FLOOR_DIV: (Op.FloorDiv, OptionsSerializer("FloorDivOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.REDUCE_ANY: (Op.Any, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
    BuiltinOperator.SQUARE: (Op.Square, OptionsSerializer("SquareOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.ZEROS_LIKE: (Op.ZerosLike, OptionsSerializer("ZerosLikeOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.FILL: (Op.Fill, OptionsSerializer("FillOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.FLOOR_MOD: (Op.FloorMod, OptionsSerializer("FloorModOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.RANGE: (Op.Range, OptionsSerializer("RangeOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.RESIZE_NEAREST_NEIGHBOR: (
        Op.ResizeNearestNeighbor,
        OptionsSerializer("ResizeNearestNeighborOptions", ("align_corners", "half_pixel_centers")),
        TFLITE_IFM_INDICES,
    ),
    BuiltinOperator.LEAKY_RELU: (Op.LeakyRelu, OptionsSerializer("LeakyReluOptions", ("alpha",)), TFLITE_IFM_INDICES),
    BuiltinOperator.SQUARED_DIFFERENCE: (
        Op.SquaredDifference,
        OptionsSerializer("SquaredDifferenceOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.MIRROR_PAD: (Op.MirrorPad, OptionsSerializer("MirrorPadOptions", ("mode",)), TFLITE_NO_INDICES),
    BuiltinOperator.ABS: (Op.Abs, OptionsSerializer("AbsOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.SPLIT_V: (Op.SplitV, OptionsSerializer("SplitVOptions", ("num_splits",)), TFLITE_IFM_INDICES),
    BuiltinOperator.UNIQUE: (
        Op.Unique,
        OptionsSerializer("UniqueOptions", (("idx_out_type", datatype_deserialize, datatype_serialize),)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.CEIL: (Op.Ceil, None, TFLITE_NO_INDICES),
    BuiltinOperator.REVERSE_V2: (Op.ReverseV2, OptionsSerializer("ReverseV2Options"), TFLITE_NO_INDICES),
    BuiltinOperator.ADD_N: (Op.AddN, OptionsSerializer("AddNOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.GATHER_ND: (Op.GatherNd, OptionsSerializer("GatherNdOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.COS: (Op.Cos, OptionsSerializer("CosOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.WHERE: (Op.Where, OptionsSerializer("WhereOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.RANK: (Op.Rank, OptionsSerializer("RankOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.ELU: (Op.Elu, None, TFLITE_NO_INDICES),
    BuiltinOperator.REVERSE_SEQUENCE: (
        Op.ReverseSequence,
        OptionsSerializer("ReverseSequenceOptions", ("batch_dim", "seq_dim")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.MATRIX_DIAG: (Op.MatrixDiag, OptionsSerializer("MatrixDiagOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.QUANTIZE: (Op.Quantize, OptionsSerializer("QuantizeOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.MATRIX_SET_DIAG: (Op.MatrixSetDiag, OptionsSerializer("MatrixSetDiagOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.ROUND: (Op.Round, None, TFLITE_NO_INDICES),
    BuiltinOperator.HARD_SWISH: (Op.HardSwish, OptionsSerializer("HardSwishOptions"), TFLITE_IFM_INDICES),
    BuiltinOperator.IF: (
        Op.If,
        OptionsSerializer("IfOptions", ("else_subgraph_index", "then_subgraph_index")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.WHILE: (
        Op.While,
        OptionsSerializer("WhileOptions", ("body_subgraph_index", "cond_subgraph_index")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.NON_MAX_SUPPRESSION_V4: (
        Op.NonMaxSuppressionV4,
        OptionsSerializer("NonMaxSuppressionV4Options"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.NON_MAX_SUPPRESSION_V5: (
        Op.NonMaxSuppressionV5,
        OptionsSerializer("NonMaxSuppressionV5Options"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.SCATTER_ND: (Op.ScatterNd, OptionsSerializer("ScatterNdOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.SELECT_V2: (Op.SelectV2, OptionsSerializer("SelectV2Options"), TFLITE_NO_INDICES),
    BuiltinOperator.DENSIFY: (Op.Densify, OptionsSerializer("DensifyOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.SEGMENT_SUM: (Op.SegmentSum, OptionsSerializer("SegmentSumOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.BATCH_MATMUL: (
        Op.BatchMatMul,
        OptionsSerializer("BatchMatMulOptions", ("adj_x", "adj_y", "asymmetric_quantize_inputs")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.CUMSUM: (
        Op.Cumsum,
        OptionsSerializer("CumsumOptions", ("exclusive", "reverse")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.CALL_ONCE: (
        Op.CallOnce,
        OptionsSerializer("CallOnceOptions", ("init_subgraph_index",)),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.BROADCAST_TO: (Op.BroadcastTo, OptionsSerializer("BroadcastToOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.RFFT2D: (Op.Rfft2D, OptionsSerializer("Rfft2dOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.CONV_3D: (
        Op.Conv3D,
        OptionsSerializer(
            "Conv3DOptions",
            (
                "dilation_d_factor",
                "dilation_h_factor",
                "dilation_w_factor",
                fused_act,
                padding,
                "stride_d",
                "stride_h",
                "stride_w",
            ),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.IMAG: (Op.Imag, None, TFLITE_NO_INDICES),
    BuiltinOperator.REAL: (Op.Real, None, TFLITE_NO_INDICES),
    BuiltinOperator.COMPLEX_ABS: (Op.ComplexAbs, None, TFLITE_NO_INDICES),
    BuiltinOperator.HASHTABLE: (
        Op.Hashtable,
        OptionsSerializer("HashtableOptions", ("key_dtype", "table_id", "value_dtype")),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.HASHTABLE_FIND: (Op.HashtableFind, OptionsSerializer("HashtableFindOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.HASHTABLE_IMPORT: (
        Op.HashtableImport,
        OptionsSerializer("HashtableImportOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.HASHTABLE_SIZE: (Op.HashtableSize, OptionsSerializer("HashtableSizeOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.REDUCE_ALL: (Op.ReduceAll, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
    BuiltinOperator.CONV_3D_TRANSPOSE: (
        Op.Conv3DTranspose,
        OptionsSerializer(
            "Conv3DOptions",
            (
                "dilation_d_factor",
                "dilation_h_factor",
                "dilation_w_factor",
                fused_act,
                padding,
                "stride_d",
                "stride_h",
                "stride_w",
            ),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.VAR_HANDLE: (
        Op.VarHandle,
        OptionsSerializer(
            "VarHandleOptions",
            (
                "container",
                "shared_name",
            ),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.READ_VARIABLE: (Op.ReadVariable, OptionsSerializer("ReadVariableOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.ASSIGN_VARIABLE: (Op.AssignVariable, OptionsSerializer("AssignVariableOptions"), TFLITE_NO_INDICES),
    BuiltinOperator.BROADCAST_ARGS: (Op.BroadcastArgs, None, TFLITE_NO_INDICES),
    BuiltinOperator.RANDOM_STANDARD_NORMAL: (
        Op.RandomStandardNormal,
        OptionsSerializer(
            "RandomOptions",
            (
                "seed",
                "seed2",
            ),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.BUCKETIZE: (
        Op.Bucketize,
        OptionsSerializer(
            "BucketizeOptions",
            ("boundaries", "boundaries_as_numpy", "boundaries_length", "boundaries_is_none"),
        ),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.RANDOM_UNIFORM: (Op.RandomUniform, None, TFLITE_NO_INDICES),
    BuiltinOperator.MULTINOMIAL: (Op.Multinomial, None, TFLITE_NO_INDICES),
    BuiltinOperator.GELU: (Op.Gelu, OptionsSerializer("GeluOptions", ("approximate",)), TFLITE_NO_INDICES),
    BuiltinOperator.DYNAMIC_UPDATE_SLICE: (
        Op.DynamicUpdateSlice,
        OptionsSerializer("DynamicUpdateSliceOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.UNSORTED_SEGMENT_PROD: (
        Op.UnsortedSegmentProd,
        OptionsSerializer("UnsortedSegmentProdOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.UNSORTED_SEGMENT_MAX: (
        Op.UnsortedSegmentMax,
        OptionsSerializer("UnsortedSegmentMaxOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.UNSORTED_SEGMENT_MIN: (
        Op.UnsortedSegmentMin,
        OptionsSerializer("UnsortedSegmentMinOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.UNSORTED_SEGMENT_SUM: (
        Op.UnsortedSegmentSum,
        OptionsSerializer("UnsortedSegmentSumOptions"),
        TFLITE_NO_INDICES,
    ),
    BuiltinOperator.CUSTOM: (Op.Custom, CustomOptionsSerializer(), TFLITE_NO_INDICES),
}

builtin_operator_inv_map = {v[0]: (k, v[1], v[2]) for k, v in builtin_operator_map.items()}

builtin_operator_inv_map[Op.CustomNpuOp] = (BuiltinOperator.CUSTOM, CustomOptionsSerializer(), TFLITE_NO_INDICES)

BUILTIN_OPERATOR_UNKNOWN = "UNKNOWN"

builtin_operator_name_map = {v: k for k, v in vars(BuiltinOperator).items()}


def optype_to_builtintype(op_type: Op):
    if op_type in builtin_operator_inv_map:
        return builtin_operator_name_map[builtin_operator_inv_map[op_type][0]]
    else:
        return BUILTIN_OPERATOR_UNKNOWN