import collections.abc
import numbers
import os
import random
import struct
import sys
from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
    TypeVar,
    Union,
    cast,
)

import flatbuffers
import numpy as np
import onnx
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from onnx import numpy_helper

# Generated by `flatc`
sys.path.append(os.path.dirname(__file__))
import FlatBuffers.Dl.AlignedBytes as AlignedBytes
import FlatBuffers.Dl.Attribute as Attribute
import FlatBuffers.Dl.AttributeF as AttributeF
import FlatBuffers.Dl.AttributeI as AttributeI
import FlatBuffers.Dl.AttributeType as AttributeType
import FlatBuffers.Dl.Dimension as Dimension
import FlatBuffers.Dl.DimensionValue as DimensionValue
import FlatBuffers.Dl.DimensionValueType as DimensionValueType
import FlatBuffers.Dl.Function as Function
import FlatBuffers.Dl.Graph as Graph
import FlatBuffers.Dl.Model as Model
import FlatBuffers.Dl.Node as Node
import FlatBuffers.Dl.OperatorSetId as OperatorSetId
import FlatBuffers.Dl.StringStringEntry as StringStringEntry
import FlatBuffers.Dl.Tensor as Tensor
import FlatBuffers.Dl.TensorDataType as TensorDataType
import FlatBuffers.Dl.TensorShape as TensorShape
import FlatBuffers.Dl.TensorTypeAndShape as TensorTypeAndShape
import FlatBuffers.Dl.TypeInfo as TypeInfo
import FlatBuffers.Dl.TypeInfoValue as TypeInfoValue
import FlatBuffers.Dl.ValueInfo as ValueInfo
import mapping as mapping

global_fbs_builder = None


def get_global_fbs_builder():
    global global_fbs_builder
    if global_fbs_builder is None:
        global_fbs_builder = flatbuffers.Builder(0)
    return global_fbs_builder


def reset_global_fbs_builder():
    global global_fbs_builder
    global_fbs_builder = None


def make_node(
    op_type: str,
    inputs: Sequence[str],
    outputs: Sequence[str],
    name: Optional[str] = None,
    doc_string: Optional[str] = None,
    domain: Optional[str] = None,
    **kwargs: Any,
) -> int:
    """Construct a Node.

    Args:
        op_type (string): The name of the operator to construct
        inputs (list of string): list of input names
        outputs (list of string): list of output names
        name (string, default None): optional unique identifier for NodeFbs
        doc_string (string, default None): optional documentation string for NodeFbs
        domain (string, default None): optional domain for NodeFbs.
            If it's None, we will just use default domain (which is empty)
        **kwargs (dict): the attributes of the node.  The acceptable values
            are documented in :func:`make_attribute`.

    Returns:
        Node fbs offset
    """
    builder = get_global_fbs_builder()
    op_type = builder.CreateString(op_type)

    inputs = [builder.CreateString(i) for i in inputs]
    Node.StartInputVector(builder, len(inputs))
    for i in reversed(inputs):
        builder.PrependUOffsetTRelative(i)
    inputs = builder.EndVector()

    outputs = [builder.CreateString(o) for o in outputs]
    Node.StartOutputVector(builder, len(outputs))
    for o in reversed(outputs):
        builder.PrependUOffsetTRelative(o)
    outputs = builder.EndVector()

    if name:
        name = builder.CreateString(name)
    if doc_string:
        doc_string = builder.CreateString(doc_string)
    if domain:
        domain = builder.CreateString(domain)

    attrs = None
    if kwargs:
        attrs = [make_attribute(key, value) for key, value in sorted(kwargs.items()) if value is not None]
        Node.StartAttributeVector(builder, len(attrs))
        for attr in reversed(attrs):
            builder.PrependUOffsetTRelative(attr)
        attrs = builder.EndVector()

    Node.Start(builder)
    Node.AddOpType(builder, op_type)
    Node.AddInput(builder, inputs)
    Node.AddOutput(builder, outputs)
    if name:
        Node.AddName(builder, name)
    if doc_string:
        Node.AddDocString(builder, doc_string)
    if domain:
        Node.AddDomain(builder, domain)
    if attrs is not None:
        Node.AddAttribute(builder, attrs)
    return Node.End(builder)


def make_graph_test_value(
    valuesForTest: Dict[str, Dict[str, np.ndarray]] = None, exponents: Dict[str, List[int]] = None
) -> Tuple[Sequence[int], Sequence[int]]:
    """Construct the test_inputs_value, test_outputs_value of Graph

    Args:
        valuesForTest (Dict[str, Dict[str, np.ndarray]]): the test values used to compare accuracy.
                                                          The input format is as follows:
            {
                'inputs': {
                    'input_0_name': np.ndarray
                    ......
                    'input_n_name': np.ndarray
                },
                'outputs': {
                    'output_0_name': np.ndarray
                    ......
                    'output_n_name': np.ndarray
                },
            }

    Returns:
        Tuple[Sequence[Tensor fbs offset], Sequence[Tensor fbs offset]]
    """
    input_test_values = []
    output_test_values = []
    if valuesForTest is not None:
        # Get the input test value
        for input_name, input_value in valuesForTest.get('inputs', {}).items():
            var_exponents = exponents.get(input_name, [0])
            tensor_type = np_dtype_to_tensor_dtype(input_value.dtype)
            input_shape = input_value.shape
            if len(input_shape) > 1:
                input_value = input_value.flatten()
            if len(input_shape) == 0:
                input_shape = [1]
            input_value_fbs = make_tensor(
                input_name, tensor_type, input_shape, input_value, raw=True, exponents=var_exponents
            )
            input_test_values.append(input_value_fbs)

        # Get the output test value
        for output_name, output_value in valuesForTest.get('outputs', {}).items():
            output_value = np.atleast_1d(output_value)
            var_exponents = exponents.get(output_name, [0])
            tensor_type = np_dtype_to_tensor_dtype(output_value.dtype)
            output_shape = output_value.shape
            if len(output_shape) > 1:
                output_value = output_value.flatten()
            output_value_fbs = make_tensor(
                output_name, tensor_type, output_shape, output_value, raw=True, exponents=var_exponents
            )
            output_test_values.append(output_value_fbs)

    return input_test_values, output_test_values


def make_graph(
    nodes: Sequence[int],
    name: str,
    inputs: Sequence[int],
    outputs: Sequence[int],
    initializer: Optional[Sequence[int]] = None,
    doc_string: Optional[str] = None,
    value_info: Optional[Sequence[int]] = None,
    test_inputs_value: Optional[Sequence[int]] = None,
    test_outputs_value: Optional[Sequence[int]] = None,
) -> int:
    """Construct a Graph

    Args:
        nodes: list of Node fbs offset
        name (string): graph name
        inputs: list of ValueInfo fbs offset
        outputs: list of ValueInfo fbs offset
        initializer: list of Tensor fbs offset
        doc_string (string): graph documentation
        value_info: list of ValueInfo fbs offset
    Returns:
        Graph fbs offset
    """
    builder = get_global_fbs_builder()

    name = builder.CreateString(name)
    if initializer is not None:
        Graph.StartInitializerVector(builder, len(initializer))
        for i in reversed(initializer):
            builder.PrependUOffsetTRelative(i)
        initializer = builder.EndVector()
    if value_info is not None:
        Graph.StartValueInfoVector(builder, len(value_info))
        for i in reversed(value_info):
            builder.PrependUOffsetTRelative(i)
        value_info = builder.EndVector()
    if test_inputs_value is not None:
        Graph.StartTestInputsValueVector(builder, len(test_inputs_value))
        for i in reversed(test_inputs_value):
            builder.PrependUOffsetTRelative(i)
        test_inputs_value = builder.EndVector()
    if test_outputs_value is not None:
        Graph.StartTestOutputsValueVector(builder, len(test_outputs_value))
        for i in reversed(test_outputs_value):
            builder.PrependUOffsetTRelative(i)
        test_outputs_value = builder.EndVector()

    Graph.StartNodeVector(builder, len(nodes))
    for i in reversed(nodes):
        builder.PrependUOffsetTRelative(i)
    nodes = builder.EndVector()

    Graph.StartInputVector(builder, len(inputs))
    for i in reversed(inputs):
        builder.PrependUOffsetTRelative(i)
    inputs = builder.EndVector()

    Graph.StartOutputVector(builder, len(outputs))
    for i in reversed(outputs):
        builder.PrependUOffsetTRelative(i)
    outputs = builder.EndVector()

    Graph.Start(builder)
    if initializer is not None:
        Graph.AddInitializer(builder, initializer)
    if value_info is not None:
        Graph.AddValueInfo(builder, value_info)
    if test_inputs_value is not None:
        Graph.AddTestInputsValue(builder, test_inputs_value)
    if test_outputs_value is not None:
        Graph.AddTestOutputsValue(builder, test_outputs_value)
    Graph.AddNode(builder, nodes)
    Graph.AddName(builder, name)
    Graph.AddInput(builder, inputs)
    Graph.AddOutput(builder, outputs)
    if doc_string:
        Graph.AddDocString(builder, builder.CreateString(doc_string))
    return Graph.End(builder)


def make_metadata_props(metadata_props: Dict[str, str]) -> Sequence[int]:
    """Construct a StringStringEntry

    Args:
        metadata_props (Dict[str, str]): the key/value of metadata_props
    Returns:
        StringStringEntry fbs offset
    """
    builder = get_global_fbs_builder()
    _metadata_props = []
    if metadata_props:
        for k, v in metadata_props.items():
            k_offset = builder.CreateString(k)
            v_offset = builder.CreateString(v)
            StringStringEntry.Start(builder)
            StringStringEntry.AddKey(builder, k_offset)
            StringStringEntry.AddValue(builder, v_offset)
            _metadata_props.append(StringStringEntry.End(builder))
    return _metadata_props


def make_model(graph: int, **kwargs: Any) -> bytes:
    """Construct a Model

    Args:
        graph (int): graph fbs offset
        **kwargs: any attribute to add to the returned instance
    Returns:
        model fbs data in bytes
    """
    builder = get_global_fbs_builder()

    opset_imports = kwargs.pop("opset_imports", None)  # type: ignore
    if opset_imports is not None:
        Model.StartOpsetImportVector(builder, len(opset_imports))
        for i in reversed(opset_imports):
            builder.PrependUOffsetTRelative(i)
        opset_imports = builder.EndVector()
    else:
        Model.StartOpsetImportVector(builder, 0)
        opset_imports = builder.EndVector()

    functions = kwargs.pop("functions", None)  # type: ignore
    if functions is not None:
        Model.StartFunctionsVector(builder, len(functions))
        for i in reversed(functions):
            builder.PrependUOffsetTRelative(i)
        functions = builder.EndVector()

    producer_name = kwargs.pop("producer_name", None)  # type: ignore
    if producer_name is not None:
        producer_name = builder.CreateString(producer_name)

    producer_version = kwargs.pop("producer_version", None)  # type: ignore
    if producer_version is not None:
        producer_version = builder.CreateString(producer_version)

    domain = kwargs.pop("domain", None)  # type: ignore
    if domain is not None:
        domain = builder.CreateString(domain)

    doc_string = kwargs.pop("doc_string", None)  # type: ignore
    if doc_string:
        doc_string = builder.CreateString(doc_string)

    metadata_props = kwargs.pop("metadata_props", None)
    if metadata_props:
        Model.StartMetadataPropsVector(builder, len(metadata_props))
        for i in reversed(metadata_props):
            builder.PrependUOffsetTRelative(i)
        metadata_props = builder.EndVector()

    Model.Start(builder)
    Model.AddGraph(builder, graph)
    Model.AddOpsetImport(builder, opset_imports)
    if functions is not None:
        Model.AddFunctions(builder, functions)
    ir_version = kwargs.pop("ir_version", None)  # type: ignore
    if ir_version is not None:
        Model.AddIrVersion(builder, ir_version)
    if producer_name:
        Model.AddProducerName(builder, producer_name)
    if producer_version:
        Model.AddProducerVersion(builder, producer_version)
    if domain:
        Model.AddDomain(builder, domain)
    model_version = kwargs.pop("model_version", None)  # type: ignore
    if model_version is not None:
        Model.AddModelVersion(builder, model_version)
    if doc_string:
        Model.AddDocString(builder, doc_string)
    if metadata_props:
        Model.AddMetadataProps(builder, metadata_props)
    model = Model.End(builder)
    builder.Finish(model)
    model_data = builder.Output()
    reset_global_fbs_builder()
    return model_data


def make_tensor(
    name: str,
    data_type: int,
    dims: Sequence[int],
    vals: Any,
    raw: bool = False,
    exponents: Sequence[int] = [0],
    doc_string: str = "",
) -> int:
    """Make a Tensor with specified arguments.  If raw is False, this
    function will choose the corresponding proto field to store the
    values based on data_type. If raw is True, use "raw_data" proto
    field to store the values, and values should be of type bytes in
    this case.

    Args:
        name (string): tensor name
        data_type (int): a value such as TensorDataType.FLOAT
        dims (List[int]): shape
        vals: values
        raw (bool): if True, vals contains the serialized content of the tensor,
            otherwise, vals should be a list of values of the type defined by *data_type*

    Returns:
        Tensor offset
    """
    builder = get_global_fbs_builder()
    name = builder.CreateString(name)
    # Check number of vals specified equals tensor size
    expected_size = 1
    if type(vals) is np.ndarray and len(vals.shape) > 1:
        vals = vals.flatten()
    for d in dims:
        expected_size *= d
    if len(vals) != expected_size:
        raise ValueError(
            f"Number of values does not match tensor's size. Expected {expected_size}, but it is {len(vals)}. "
        )
    dims = builder.CreateNumpyVector(np.array(dims, dtype=np.int64))
    exponents = builder.CreateNumpyVector(np.array(exponents, dtype=np.int64))
    if doc_string:
        doc_string = builder.CreateString(doc_string)

    if data_type == TensorDataType.TensorDataType().STRING and raw:
        raise TypeError("Can not use raw_data to store string type.")

    np_dtype = tensor_dtype_to_np_dtype(data_type)

    if raw:
        n = 16 // np_dtype.itemsize
        if len(vals) % n != 0:
            padding = n - len(vals) % n
            vals = np.pad(vals, (0, padding), mode='constant', constant_values=0)
        vals = np.reshape(vals, [-1, n])
        vals = vals.astype(np_dtype)
        Tensor.StartRawDataVector(builder, len(vals))
        for val in reversed(vals):
            AlignedBytes.CreateAlignedBytes(builder, val.tobytes())
        vals = builder.EndVector()
    else:
        if data_type == TensorDataType.TensorDataType().FLOAT16:
            raise TypeError("Don't support FLOAT16.")
        elif data_type == TensorDataType.TensorDataType().BOOL:
            vals = np.array(vals).astype(int)
        elif data_type == TensorDataType.TensorDataType().STRING:
            vals = np.array(vals).astype(bytes)
        else:
            vals = np.array(vals)
        vals = builder.CreateNumpyVector(vals)

    Tensor.Start(builder)
    Tensor.AddDataType(builder, data_type)
    Tensor.AddName(builder, name)
    if doc_string:
        Tensor.AddDocString(builder, doc_string)
    if raw:
        Tensor.AddRawData(builder, vals)
    else:
        add_field = tensor_dtype_to_field(data_type)
        add_field(builder, vals)
    Tensor.AddDims(builder, dims)
    Tensor.AddExponents(builder, exponents)
    return Tensor.End(builder)


def _to_bytes(value: Union[str, bytes]) -> bytes:
    """Coerce a string (or bytes) value into UTF-8 bytes."""
    return value if isinstance(value, bytes) else value.encode("utf-8")


def _to_str(value: Union[str, bytes]) -> str:
    """Coerce a string (or UTF-8 bytes) value into str."""
    return value if isinstance(value, str) else value.decode("utf-8", errors="ignore")


def make_attribute(
    key: str,
    value: Any,
    doc_string: Optional[str] = None,
    attr_type: Optional[int] = None,
) -> int:
    """Makes an Attribute based on the value type."""
    builder = get_global_fbs_builder()

    key = builder.CreateString(key)
    if doc_string:
        doc_string = builder.CreateString(doc_string)

    attr_i = None
    attr_f = None
    attr_s = None
    attr_t = None
    attr_ints = None
    attr_floats = None
    attr_strings = None
    # Singular cases
    if isinstance(value, numbers.Integral):
        attr_i = int(value)
    elif isinstance(value, numbers.Real):
        attr_f = float(value)
    elif isinstance(value, (str, bytes)):
        # Encode strings into utf-8
        value = _to_bytes(value)
        attr_s = builder.CreateByteVector(value)
    elif isinstance(value, onnx.TensorProto):
        value_data = numpy_helper.to_array(value)
        if (
            value.data_type >= TensorDataType.TensorDataType.UNDEFINED
            and value.data_type <= TensorDataType.TensorDataType.UINT64
        ):
            attr_t = make_tensor(
                name=value.name,
                data_type=value.data_type,
                dims=value.dims,
                vals=value_data,
                raw=True,
                doc_string=value.doc_string,
            )
        else:
            raise ValueError(f"Don't support type {onnx.TensorProto.DataType.Name(value.data_type)}")

    # Iterable cases
    elif isinstance(value, collections.abc.Iterable):
        value = list(value)
        if len(value) == 0 and attr_type is None:
            raise ValueError(f"Could not infer attribute `{key}` type from empty iterator")
        if attr_type is None:
            types = {type(v) for v in value}
            for exp_t, exp_enum in (
                (numbers.Integral, AttributeType.AttributeType().INTS),
                (numbers.Real, AttributeType.AttributeType().FLOATS),
                ((str, bytes), AttributeType.AttributeType().STRINGS),
            ):
                if all(issubclass(t, exp_t) for t in types):  # type: ignore[arg-type]
                    attr_type = exp_enum
                    break
            if attr_type is None:
                raise ValueError("Could not infer the attribute type from the elements of the passed Iterable value.")

        if attr_type == AttributeType.AttributeType().INTS:
            value = np.array(value, dtype=np.int64)
            attr_ints = builder.CreateNumpyVector(value)
        elif attr_type == AttributeType.AttributeType().FLOATS:
            value = np.array(value).astype(float)
            attr_floats = builder.CreateNumpyVector(value)
        elif attr_type == AttributeType.AttributeType().STRINGS:
            attr_strings = [builder.CreateString(s) for s in value]
            Attribute.StartStringsVector(builder, len(attr_strings))
            for s in reversed(attr_strings):
                builder.PrependUOffsetTRelative(s)
            attr_strings = builder.EndVector()
        else:
            raise AssertionError()  # Should not reach since `ValueError` must be raised in attr_type checking
    else:
        raise TypeError(f"'{value}' is not an accepted attribute value.")

    Attribute.AttributeStart(builder)
    Attribute.AddName(builder, key)
    if doc_string:
        Attribute.AddDocString(builder, doc_string)

    if attr_i is not None:
        Attribute.AddI(builder, AttributeI.CreateAttributeI(builder, attr_i))
        Attribute.AddAttrType(builder, AttributeType.AttributeType().INT)
    if attr_f is not None:
        Attribute.AddF(builder, AttributeF.CreateAttributeF(builder, attr_f))
        Attribute.AddAttrType(builder, AttributeType.AttributeType().FLOAT)
    if attr_s is not None:
        Attribute.AddS(builder, attr_s)
        Attribute.AddAttrType(builder, AttributeType.AttributeType().STRING)
    if attr_t is not None:
        Attribute.AddT(builder, attr_t)
        Attribute.AddAttrType(builder, AttributeType.AttributeType().TENSOR)
    if attr_ints is not None:
        Attribute.AddInts(builder, attr_ints)
        Attribute.AddAttrType(builder, AttributeType.AttributeType().INTS)
    if attr_floats is not None:
        Attribute.AddFloats(builder, attr_floats)
        Attribute.AddAttrType(builder, AttributeType.AttributeType().FLOATS)
    if attr_strings is not None:
        Attribute.AddStrings(builder, attr_strings)
        Attribute.AddAttrType(builder, AttributeType.AttributeType().STRINGS)
    return Attribute.End(builder)


def make_tensor_type_info(
    elem_type: int,
    shape: Optional[Sequence[Union[str, int, None]]],
    shape_denotation: Optional[List[str]] = None,
) -> int:
    builder = get_global_fbs_builder()
    tensor_shape = None
    if shape is not None:
        if shape_denotation and len(shape_denotation) != len(shape):
            raise ValueError("Invalid shape_denotation. Must be of the same length as shape.")
        dim_value = []
        for i, d in enumerate(shape):
            DimensionValue.Start(builder)
            if d is None:
                pass
            elif isinstance(d, int):
                DimensionValue.AddDimType(builder, DimensionValueType.DimensionValueType().VALUE)
                DimensionValue.AddDimValue(builder, d)

            elif isinstance(d, str):
                DimensionValue.AddDimType(builder, DimensionValueType.DimensionValueType().PARAM)
                DimensionValue.AddDimParam(builder, d)
            else:
                raise ValueError(f"Invalid item in shape: {d}. Needs to be of int or str.")
            dim_value.append(DimensionValue.End(builder))

        if shape_denotation is not None:
            shape_denotation = [builder.CreateString(builder, s) for s in shape_denotation]
        dim = []
        for i in range(len(shape)):
            Dimension.Start(builder)
            Dimension.AddValue(builder, dim_value[i])
            if shape_denotation is not None:
                Dimension.AddDenotation(builder, shape_denotation[i])
            dim.append(Dimension.End(builder))

        TensorShape.TensorShapeStartDimVector(builder, len(shape))
        for d in reversed(dim):
            builder.PrependUOffsetTRelative(d)
        dim = builder.EndVector()
        TensorShape.TensorShapeStart(builder)
        TensorShape.AddDim(builder, dim)
        tensor_shape = TensorShape.End(builder)

    TensorTypeAndShape.Start(builder)
    if tensor_shape is not None:
        TensorTypeAndShape.AddShape(builder, tensor_shape)
    TensorTypeAndShape.AddElemType(builder, elem_type)
    tensor_type_and_shape = TensorTypeAndShape.End(builder)

    TypeInfo.Start(builder)
    TypeInfo.AddValueType(builder, TypeInfoValue.TypeInfoValue().tensor_type)
    TypeInfo.AddValue(builder, tensor_type_and_shape)
    return TypeInfo.End(builder)


def make_tensor_value_info(
    name: str,
    elem_type: int,
    shape: Optional[Sequence[Union[str, int, None]]],
    doc_string: str = "",
    shape_denotation: Optional[List[str]] = None,
    exponents: Sequence[int] = [0],
) -> int:
    """Makes a ValueInfo fbs offset based on the data type and shape."""
    builder = get_global_fbs_builder()
    name = builder.CreateString(name)
    if doc_string:
        doc_string = builder.CreateString(doc_string)
    type_info = make_tensor_type_info(elem_type, shape, shape_denotation)
    exponents = builder.CreateNumpyVector(np.array(exponents, dtype=np.int64))
    ValueInfo.Start(builder)
    ValueInfo.AddName(builder, name)
    if doc_string:
        ValueInfo.AddDocString(builder, doc_string)
    ValueInfo.AddValueInfoType(builder, type_info)
    ValueInfo.AddExponents(builder, exponents)
    return ValueInfo.End(builder)


def aes_encryptor(data: bytes):
    key = bytearray(random.getrandbits(8) for _ in range(16))
    iv = bytearray(range(16, 32, 1))
    cipher = Cipher(algorithms.AES(key), modes.CTR(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    ct_data = encryptor.update(data) + encryptor.finalize()

    hex_key = ''
    for _, byte in enumerate(key):
        hex_key += "0x{:02x}, ".format(byte)

    return ct_data, hex_key


def aes_decryptor(data: bytes, key: bytearray):
    iv = bytearray(range(16, 32, 1))
    cipher = Cipher(algorithms.AES(key), modes.CTR(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    decrypted_data = decryptor.update(data) + decryptor.finalize()
    return decrypted_data


def save(data: bytes, file_path: str, encrypt_data: bool = False):
    key = None
    if encrypt_data:
        data, key = aes_encryptor(data)

    # file header 16 bytes
    header = bytearray("EDL2", 'utf-8')
    header += struct.pack("I", 1 if encrypt_data else 0)
    header += struct.pack("I", len(data))
    header += struct.pack("x") * 4
    with open(file_path, mode="wb") as f:
        f.write(header + data)
    return key


def _sanitize_str(s: Union[str, bytes]) -> str:
    if isinstance(s, str):
        sanitized = s
    elif isinstance(s, bytes):
        sanitized = s.decode("utf-8", errors="ignore")
    else:
        sanitized = str(s)
    if len(sanitized) < 64:  # noqa: PLR2004
        return sanitized
    return sanitized[:64] + f"...<+len={(len(sanitized) - 64)}>"


def printable_attribute(
    attr: Attribute.Attribute, subgraphs: bool = False
) -> Union[str, Tuple[str, List[Graph.Graph]]]:
    content = []
    content.append(_to_str(attr.Name()))
    content.append("=")

    def str_float(f: float) -> str:
        # NB: Different Python versions print different numbers of trailing
        # decimals, specifying this explicitly keeps it consistent for all
        # versions
        return f"{f:.15g}"

    def str_int(i: int) -> str:
        return str(i)

    _T = TypeVar("_T")

    def str_list(str_elem: Callable[[_T], str], xs: Sequence[_T]) -> str:
        return "[" + ", ".join(map(str_elem, xs)) + "]"

    # for now, this logic should continue to work as long as we are running on a flatbuffers
    # implementation. If/when we switch to flatbuffers, we will need to use attr.type

    # To support printing subgraphs, if we find a graph attribute, print out
    # its name here and pass the graph itself up to the caller for later
    # printing.

    attr_type = attr.AttrType()

    graphs = []
    if attr_type == AttributeType.AttributeType.INT:
        content.append(str_int(attr.I().I()))
    elif attr_type == AttributeType.AttributeType.FLOAT:
        content.append(str_float(attr.F().F()))
    elif attr_type == AttributeType.AttributeType.STRING:
        # TODO: Bit nervous about Python 2 / Python 3 determinism implications
        content.append(repr(_sanitize_str(attr.SAsNumpy().tobytes())))
    elif attr_type == AttributeType.AttributeType.TENSOR:
        if attr.T().DimsLength() > 0:
            content.append("<Tensor>")
        else:
            # special case to print scalars
            field = tensor_dtype_to_field(attr.t().DataType())
            content.append(f"<Scalar Tensor {getattr(attr.t(), field)}>")
    elif attr_type == AttributeType.AttributeType.FLOATS:
        content.append(str_list(str_float, [attr.Floats(i) for i in range(attr.FloatsLength())]))
    elif attr_type == AttributeType.AttributeType.INTS:
        content.append(str_list(str_int, [attr.Ints(i) for i in range(attr.IntsLength())]))
    elif attr_type == AttributeType.AttributeType.STRINGS:
        # TODO: Bit nervous about Python 2 / Python 3 determinism implications
        content.append(str(list(map(_sanitize_str, [attr.Strings(i) for i in range(attr.StringsLength())]))))
    elif attr_type == AttributeType.AttributeType.TENSORS:
        content.append("[<Tensor>, ...]")
    else:
        content.append("<Unknown>")
    if subgraphs:
        return " ".join(content), graphs
    return " ".join(content)


def printable_dim(dim: Dimension.Dimension) -> str:
    if dim.Value().DimType() == DimensionValueType.DimensionValueType.UNKNOWN:
        return "?"
    elif dim.Value().DimType() == DimensionValueType.DimensionValueType.VALUE:
        which = dim.Value().DimValue()
    elif dim.Value().DimType() == DimensionValueType.DimensionValueType.PARAM:
        which = dim.Value().DimParam()
    return str(which)


def get_class_attribute_name(obj: Any, value: int) -> str:
    ret: str = ""
    # 获取类和实例的所有属性名字
    all_attributes = dir(obj)
    filtered_attributes = [attr for attr in all_attributes if not attr.startswith('__') and not attr.endswith('__')]
    for attr in filtered_attributes:
        if getattr(obj, attr) == value:
            ret = attr
            break

    return ret


def printable_type(t: TypeInfo.TypeInfo) -> str:
    if t.ValueType() == TypeInfoValue.TypeInfoValue.tensor_type:
        tensor_type_and_shape = TensorTypeAndShape.TensorTypeAndShape()
        tensor_type_and_shape.Init(t.Value().Bytes, t.Value().Pos)
        s = get_class_attribute_name(TensorDataType.TensorDataType(), tensor_type_and_shape.ElemType())
        if tensor_type_and_shape.Shape() is not None:
            if tensor_type_and_shape.Shape().DimLength():
                s += str(
                    ", "
                    + "x".join(
                        map(
                            printable_dim,
                            [
                                tensor_type_and_shape.Shape().Dim(i)
                                for i in range(tensor_type_and_shape.Shape().DimLength())
                            ],
                        )
                    )
                )
            else:
                s += ", scalar"
        return s  # type: ignore[no-any-return]
    if t.ValueType() == TypeInfoValue.TypeInfoValue.NONE:
        return ""
    return f"Unknown type {t.valueType}"


def printable_value_info(v: ValueInfo.ValueInfo) -> str:
    s = f"%{_to_str(v.Name())}"
    if v.ValueInfoType():
        s = f"{s}[{printable_type(v.ValueInfoType())}]"
    if v.ExponentsLength() > 0:
        s += f", exponents: {repr(v.ExponentsAsNumpy().tolist())}"
    return s


def printable_tensor(t: Tensor.Tensor) -> str:
    s = f"%{_to_str(t.Name())}["
    s += get_class_attribute_name(TensorDataType.TensorDataType(), t.DataType())
    if t.DimsLength():
        s += str(", " + "x".join(map(str, [t.Dims(i) for i in range(t.DimsLength())])))
    else:
        s += ", scalar"
    s += "]"
    return s


def printable_tensor_value(t: Tensor.Tensor) -> str:
    s = f"%{_to_str(t.Name())}, "
    if not t.DimsIsNone():
        s += f"shape: {repr(t.DimsAsNumpy().tolist())}, "
    if not t.ExponentsIsNone():
        s += f"exponents: {repr(t.ExponentsAsNumpy().tolist())}, "
    if t.DocString() is not None:
        s += f"docString: {t.DocString()}, "

    if not t.FloatDataIsNone():
        s += f"value: {repr(t.FloatDataAsNumpy().tolist())}"
    elif not t.Int32DataIsNone():
        s += f"value: {repr(t.Int32DataAsNumpy().tolist())}"
    elif not t.StringDataIsNone():
        string_list = [t.StringData(i) for i in range(t.StringDataLength())]
        s += f"value: {repr(string_list)}"
    elif not t.Int64DataIsNone():
        s += f"value: {repr(t.Int64DataAsNumpy().tolist())}"
    elif not t.RawDataIsNone():
        np.set_printoptions(threshold=np.inf)
        chunks = []
        for i in range(t.RawDataLength()):
            chunk = t.RawData(i)
            chunks.append(np.frombuffer(chunk.BytesAsNumpy().tobytes(), dtype=tensor_dtype_to_np_dtype(t.DataType())))
        chunks = np.array(chunks).flatten()
        s += f"\nvalue: {repr(chunks)}"
    elif not t.DoubleDataIsNone():
        s += f"value: {repr(t.DoubleDataAsNumpy().tolist())}"
    elif not t.Uint64DataIsNone():
        s += f"value: {repr(t.Uint64DataAsNumpy().tolist())}"

    return s


def printable_node(
    node: Node.Node, prefix: str = "", subgraphs: bool = False
) -> Union[str, Tuple[str, List[Graph.Graph]]]:
    content = []
    if node.OutputLength():
        content.append(", ".join([f"%{_to_str(node.Output(i))}" for i in range(node.OutputLength())]))
        content.append("=")
    # To deal with nested graphs
    graphs: List[Graph.Graph] = []
    printed_attrs = []
    if node.AttributeLength():
        for i in range(node.AttributeLength()):
            attr = node.Attribute(i)
            if subgraphs:
                printed_attr_subgraphs = printable_attribute(attr, subgraphs)
                if not isinstance(printed_attr_subgraphs[1], list):
                    raise TypeError(f"printed_attr_subgraphs[1] must be an instance of {list}.")
                graphs.extend(printed_attr_subgraphs[1])
                printed_attrs.append(printed_attr_subgraphs[0])
            else:
                printed = printable_attribute(attr)
                if not isinstance(printed, str):
                    raise TypeError(f"printed must be an instance of {str}.")
                printed_attrs.append(printed)
    printed_attributes = ", ".join(sorted(printed_attrs))
    printed_inputs = ", ".join([f"%{_to_str(node.Input(i))}" for i in range(node.InputLength())])
    if node.AttributeLength():
        content.append(f"{_to_str(node.OpType())}[{printed_attributes}]({printed_inputs})")
    else:
        content.append(f"{_to_str(node.OpType())}({printed_inputs})")
    if subgraphs:
        return prefix + " ".join(content), graphs
    return prefix + " ".join(content)


def printable_model_properties(model: Model.Model) -> str:
    s = ""
    s += "metadata_props:\n"
    for i in range(model.MetadataPropsLength()):
        metadata_prop: StringStringEntry.StringStringEntry = model.MetadataProps(i)
        k = metadata_prop.Key()
        v = metadata_prop.Value()
        s += f"{_to_str(k)}: {_to_str(v)}\n"
    return s


def printable_graph(
    model: bytes,
    prefix: str = "",
    print_initializer_value: bool = False,
    print_value_info: bool = False,
    print_test_value: bool = False,
) -> str:
    """Display a Graph as a string.

    Args:
        graph (Graph.GraphT): the graph to display
        prefix (string): prefix of every line

    Returns:
        string
    """
    model = Model.Model.GetRootAs(model, 0)
    graph = model.Graph()

    content = []
    content.append("model properties:\n" + printable_model_properties(model) + "\n")

    indent = prefix + "  "
    # header
    header = ["graph", _to_str(graph.Name())]
    initializers = {graph.Initializer(i).Name() for i in range(graph.InitializerLength())}
    if graph.InputLength():
        header.append("(")
        in_strs = []  # required inputs
        in_with_init_strs = []  # optional inputs with initializer providing default value
        for i in range(graph.InputLength()):
            inp = graph.Input(i)
            if inp.Name() not in initializers:
                in_strs.append(printable_value_info(inp))
            else:
                in_with_init_strs.append(printable_value_info(inp))
        if in_strs:
            content.append(prefix + " ".join(header))
            header = []
            for line in in_strs:
                content.append(prefix + "  " + line)
        header.append(")")

        if in_with_init_strs:
            header.append("optional inputs with matching initializers (")
            content.append(prefix + " ".join(header))
            header = []
            for line in in_with_init_strs:
                content.append(prefix + "  " + line)
            header.append(")")

        # from IR 4 onwards an initializer is not required to have a matching graph input
        # so output the name, type and shape of those as well
        if len(in_with_init_strs) < len(initializers):
            graph_inputs = {graph.Input(i).Name() for i in range(graph.InputLength())}
            init_strs = [
                printable_tensor(graph.Initializer(i))
                for i in range(graph.InitializerLength())
                if graph.Initializer(i).Name() not in graph_inputs
            ]
            header.append("initializers (")
            content.append(prefix + " ".join(header))
            header = []
            for line in init_strs:
                content.append(prefix + "  " + line)
            header.append(")")

    header.append("{")
    content.append(prefix + " ".join(header))
    graphs: List[Graph.Graph] = []
    # body
    for i in range(graph.NodeLength()):
        node = graph.Node(i)
        contents_subgraphs = printable_node(node, indent, subgraphs=True)
        if not isinstance(contents_subgraphs[1], list):
            raise TypeError(f"contents_subgraphs[1] must be an instance of {list}.")
        content.append(contents_subgraphs[0])
        graphs.extend(contents_subgraphs[1])
    # tail
    tail = ["return"]
    if graph.OutputLength():
        tail.append(", ".join([f"%{_to_str(graph.Output(i).Name())}" for i in range(graph.OutputLength())]))
    content.append(indent + " ".join(tail))
    # closing bracket
    content.append(prefix + "}")
    for g in graphs:
        content.append("\n" + printable_graph(g))

    if print_initializer_value:
        initializers_value = [printable_tensor_value(graph.Initializer(i)) for i in range(graph.InitializerLength())]
        content.append("\n\n" + "initializers value:\n" + "\n".join(initializers_value))

    if print_value_info:
        values_info = [printable_value_info(graph.ValueInfo(i)) for i in range(graph.ValueInfoLength())]
        content.append("\n\n" + "values info:\n" + "\n".join(values_info))

    if print_test_value:
        test_inputs_value = [
            printable_tensor_value(graph.TestInputsValue(i)) for i in range(graph.TestInputsValueLength())
        ]
        content.append("\n\n" + "test inputs value:\n" + "\n".join(test_inputs_value))
        test_outputs_value = [
            printable_tensor_value(graph.TestOutputsValue(i)) for i in range(graph.TestOutputsValueLength())
        ]
        content.append("\n\n" + "test outputs value:\n" + "\n".join(test_outputs_value))

    return "\n".join(content)


# Following functions are used for mapping
def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
    """Convert a TensorDataType to corresponding numpy dtype. It can be used while making tensor.

    Args:
        tensor_dtype: TensorDataType

    Returns:
        numpy's data_type
    """
    return mapping.TENSOR_TYPE_MAP[tensor_dtype].np_dtype


def tensor_dtype_to_field(tensor_dtype: int) -> Callable[..., int]:
    """Convert a TensorDataType to corresponding field name for storage. It can be used while making tensors.

    Args:
        tensor_dtype: TensorDataType

    Returns:
        field name
    """
    return mapping._STORAGE_TENSOR_TYPE_TO_FIELD[mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype]


def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
    """Convert a numpy's dtype to corresponding TensorDataType. It can be used while converting numpy arrays to tensors.

    Args:
        np_dtype: numpy's data_type

    Returns:
        TensorDataType
    """
    tensor_type = mapping._NP_TYPE_TO_TENSOR_TYPE.get(np_dtype)
    if tensor_type == None:
        return cast(int, TensorDataType.TensorDataType().UNDEFINED)
    else:
        return cast(int, tensor_type)
    # return cast(
    #     int,
    #     mapping._NP_TYPE_TO_TENSOR_TYPE[np_dtype],
    # )
