onnx-rs 0.1.3

/// Top-level ONNX model container.
///
/// Corresponds to the `ModelProto` message in the ONNX protobuf schema.
/// Contains model metadata, operator set versions, and the computation
/// [`Graph`].
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::*;
///
/// let model = Model {
///     ir_version: 9,
///     producer_name: "my-framework",
///     producer_version: "1.0",
///     opset_import: vec![OperatorSetId { domain: "", version: 19 }],
///     graph: Some(Graph {
///         name: "main_graph",
///         node: vec![Node {
///             op_type: OpType::Add,
///             input: vec!["A", "B"],
///             output: vec!["C"],
///             ..Default::default()
///         }],
///         ..Default::default()
///     }),
///     ..Default::default()
/// };
///
/// assert_eq!(model.ir_version, 9);
/// assert_eq!(model.graph.as_ref().unwrap().name, "main_graph");
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Model<'a> {
    pub ir_version: i64,
    pub opset_import: Vec<OperatorSetId<'a>>,
    pub producer_name: &'a str,
    pub producer_version: &'a str,
    pub domain: &'a str,
    pub model_version: i64,
    pub doc_string: &'a str,
    pub graph: Option<Graph<'a>>,
    pub metadata_props: Vec<StringStringEntry<'a>>,
    pub training_info: Vec<TrainingInfo<'a>>,
    pub functions: Vec<Function<'a>>,
}

/// An operator set version imported by a model.
///
/// Corresponds to `OperatorSetIdProto` in the ONNX protobuf schema.
/// An empty `domain` refers to the default ONNX operator set.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::OperatorSetId;
///
/// let opset = OperatorSetId { domain: "", version: 19 };
/// assert_eq!(opset.version, 19);
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct OperatorSetId<'a> {
    pub domain: &'a str,
    pub version: i64,
}

/// A key-value string pair used for metadata properties.
///
/// Corresponds to `StringStringEntryProto` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct StringStringEntry<'a> {
    pub key: &'a str,
    pub value: &'a str,
}

/// A computation graph containing nodes, inputs, outputs, and initializers.
///
/// Corresponds to `GraphProto` in the ONNX protobuf schema.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::*;
///
/// let graph = Graph {
///     name: "my_graph",
///     node: vec![
///         Node {
///             op_type: OpType::Relu,
///             input: vec!["X"],
///             output: vec!["Y"],
///             ..Default::default()
///         },
///     ],
///     input: vec![ValueInfo { name: "X", ..Default::default() }],
///     output: vec![ValueInfo { name: "Y", ..Default::default() }],
///     ..Default::default()
/// };
///
/// assert_eq!(graph.node.len(), 1);
/// assert_eq!(graph.node[0].op_type, OpType::Relu);
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Graph<'a> {
    pub node: Vec<Node<'a>>,
    pub name: &'a str,
    pub initializer: Vec<TensorProto<'a>>,
    pub sparse_initializer: Vec<SparseTensor<'a>>,
    pub doc_string: &'a str,
    pub input: Vec<ValueInfo<'a>>,
    pub output: Vec<ValueInfo<'a>>,
    pub value_info: Vec<ValueInfo<'a>>,
    pub quantization_annotation: Vec<TensorAnnotation<'a>>,
    pub metadata_props: Vec<StringStringEntry<'a>>,
}

/// A single operation in the computation graph.
///
/// Corresponds to `NodeProto` in the ONNX protobuf schema. Each node has an
/// [`OpType`], consumes named inputs, produces named outputs, and may carry
/// [`Attribute`]s that parameterize the operation.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::*;
///
/// let node = Node {
///     name: "conv_0",
///     op_type: OpType::Conv,
///     input: vec!["X", "W"],
///     output: vec!["Y"],
///     attribute: vec![Attribute {
///         name: "kernel_shape",
///         r#type: AttributeType::Ints,
///         ints: vec![3, 3],
///         ..Default::default()
///     }],
///     ..Default::default()
/// };
///
/// assert_eq!(node.op_type, OpType::Conv);
/// assert_eq!(node.attribute[0].ints, vec![3, 3]);
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Node<'a> {
    pub input: Vec<&'a str>,
    pub output: Vec<&'a str>,
    pub name: &'a str,
    pub op_type: OpType<'a>,
    pub domain: &'a str,
    pub overload: &'a str,
    pub attribute: Vec<Attribute<'a>>,
    pub doc_string: &'a str,
    pub metadata_props: Vec<StringStringEntry<'a>>,
}

/// ONNX operator type.
///
/// Contains variants for all standard ONNX operators. Unrecognized operator
/// names are stored as [`OpType::Custom`].
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::OpType;
///
/// // Convert from a string
/// let op = OpType::from("MatMul");
/// assert_eq!(op, OpType::MatMul);
///
/// // Convert back to a string
/// assert_eq!(op.as_str(), "MatMul");
///
/// // Unknown operators become Custom
/// let custom = OpType::from("MyCustomOp");
/// assert_eq!(custom, OpType::Custom("MyCustomOp"));
/// assert_eq!(custom.as_str(), "MyCustomOp");
///
/// // Display works as expected
/// assert_eq!(format!("{}", OpType::Relu), "Relu");
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum OpType<'a> {
    // Activation
    Relu,
    LeakyRelu,
    Sigmoid,
    Tanh,
    Softmax,
    LogSoftmax,
    Elu,
    Selu,
    PRelu,
    Gelu,
    HardSigmoid,
    HardSwish,
    Softplus,
    Softsign,
    ThresholdedRelu,
    Celu,
    Mish,

    // Math (element-wise)
    Add,
    Sub,
    Mul,
    Div,
    Neg,
    Abs,
    Sqrt,
    Exp,
    Log,
    Pow,
    Ceil,
    Floor,
    Round,
    Clip,
    Sign,
    Reciprocal,
    Mod,
    BitShift,
    BitwiseAnd,
    BitwiseOr,
    BitwiseXor,
    BitwiseNot,
    Min,
    Max,
    Mean,
    Sum,

    // Trigonometric
    Sin,
    Cos,
    Tan,
    Asin,
    Acos,
    Atan,
    Sinh,
    Cosh,
    Asinh,
    Acosh,
    Atanh,

    // Convolution and pooling
    Conv,
    ConvTranspose,
    ConvInteger,
    AveragePool,
    MaxPool,
    GlobalAveragePool,
    GlobalMaxPool,
    GlobalLpPool,
    LpPool,
    MaxRoiPool,
    MaxUnpool,

    // Normalization
    BatchNormalization,
    InstanceNormalization,
    LayerNormalization,
    GroupNormalization,
    LpNormalization,
    MeanVarianceNormalization,

    // Linear algebra
    MatMul,
    MatMulInteger,
    Gemm,

    // Tensor manipulation
    Reshape,
    Transpose,
    Flatten,
    Squeeze,
    Unsqueeze,
    Expand,
    Tile,
    Pad,
    Slice,
    Split,
    SplitToSequence,
    Concat,
    ConcatFromSequence,
    DepthToSpace,
    SpaceToDepth,
    ReverseSequence,

    // Gather / Scatter
    Gather,
    GatherElements,
    GatherND,
    Scatter,
    ScatterElements,
    ScatterND,

    // Reduction
    ReduceMax,
    ReduceMin,
    ReduceMean,
    ReduceSum,
    ReduceProd,
    ReduceL1,
    ReduceL2,
    ReduceLogSum,
    ReduceLogSumExp,
    ReduceSumSquare,
    ArgMax,
    ArgMin,

    // Comparison and logic
    Equal,
    Greater,
    GreaterOrEqual,
    Less,
    LessOrEqual,
    Not,
    And,
    Or,
    Xor,
    Where,
    IsNaN,
    IsInf,

    // Shape and type
    Shape,
    Size,
    Cast,
    CastLike,
    ConstantOfShape,
    Range,
    OneHot,
    NonZero,
    TopK,
    Unique,
    EyeLike,
    Compress,

    // RNN
    LSTM,
    GRU,
    RNN,

    // Constants
    Constant,
    Identity,

    // Control flow
    If,
    Loop,
    Scan,

    // Resize
    Resize,
    Upsample,

    // Regularization
    Dropout,

    // Quantization
    QuantizeLinear,
    DequantizeLinear,
    DynamicQuantizeLinear,
    QLinearConv,
    QLinearMatMul,

    // Attention / Transformer
    Attention,
    Einsum,

    // Sequence
    SequenceConstruct,
    SequenceAt,
    SequenceEmpty,
    SequenceInsert,
    SequenceErase,
    SequenceLength,
    SequenceMap,

    // Optional
    Optional,
    OptionalGetElement,
    OptionalHasElement,

    // Image
    ImageDecoder,
    GridSample,
    RoiAlign,
    AffineGrid,

    // Misc
    CumSum,
    Det,
    Erf,
    Hardmax,
    Shrink,
    StringNormalizer,
    TfIdfVectorizer,
    NegativeLogLikelihoodLoss,
    SoftmaxCrossEntropyLoss,
    Bernoulli,
    RandomNormal,
    RandomNormalLike,
    RandomUniform,
    RandomUniformLike,
    Multinomial,
    CenterCropPad,
    Col2Im,
    Trilu,
    BlackmanWindow,
    HannWindow,
    HammingWindow,
    DFT,
    MelWeightMatrix,
    STFT,
    RegexFullMatch,
    StringConcat,
    StringSplit,

    Custom(&'a str),
}

impl OpType<'_> {
    /// Returns the ONNX string name of this operator.
    ///
    /// For standard operators this returns the canonical name (e.g. `"Relu"`).
    /// For [`OpType::Custom`] it returns the stored string.
    ///
    /// ```
    /// use onnx_rs::ast::OpType;
    ///
    /// assert_eq!(OpType::Softmax.as_str(), "Softmax");
    /// assert_eq!(OpType::Custom("FusedGelu").as_str(), "FusedGelu");
    /// ```
    /// Returns a slice of all standard ONNX operator types (excludes [`OpType::Custom`]).
    pub fn all() -> &'static [OpType<'static>] {
        &[
            OpType::Relu,
            OpType::LeakyRelu,
            OpType::Sigmoid,
            OpType::Tanh,
            OpType::Softmax,
            OpType::LogSoftmax,
            OpType::Elu,
            OpType::Selu,
            OpType::PRelu,
            OpType::Gelu,
            OpType::HardSigmoid,
            OpType::HardSwish,
            OpType::Softplus,
            OpType::Softsign,
            OpType::ThresholdedRelu,
            OpType::Celu,
            OpType::Mish,
            OpType::Add,
            OpType::Sub,
            OpType::Mul,
            OpType::Div,
            OpType::Neg,
            OpType::Abs,
            OpType::Sqrt,
            OpType::Exp,
            OpType::Log,
            OpType::Pow,
            OpType::Ceil,
            OpType::Floor,
            OpType::Round,
            OpType::Clip,
            OpType::Sign,
            OpType::Reciprocal,
            OpType::Mod,
            OpType::BitShift,
            OpType::BitwiseAnd,
            OpType::BitwiseOr,
            OpType::BitwiseXor,
            OpType::BitwiseNot,
            OpType::Min,
            OpType::Max,
            OpType::Mean,
            OpType::Sum,
            OpType::Sin,
            OpType::Cos,
            OpType::Tan,
            OpType::Asin,
            OpType::Acos,
            OpType::Atan,
            OpType::Sinh,
            OpType::Cosh,
            OpType::Asinh,
            OpType::Acosh,
            OpType::Atanh,
            OpType::Conv,
            OpType::ConvTranspose,
            OpType::ConvInteger,
            OpType::AveragePool,
            OpType::MaxPool,
            OpType::GlobalAveragePool,
            OpType::GlobalMaxPool,
            OpType::GlobalLpPool,
            OpType::LpPool,
            OpType::MaxRoiPool,
            OpType::MaxUnpool,
            OpType::BatchNormalization,
            OpType::InstanceNormalization,
            OpType::LayerNormalization,
            OpType::GroupNormalization,
            OpType::LpNormalization,
            OpType::MeanVarianceNormalization,
            OpType::MatMul,
            OpType::MatMulInteger,
            OpType::Gemm,
            OpType::Reshape,
            OpType::Transpose,
            OpType::Flatten,
            OpType::Squeeze,
            OpType::Unsqueeze,
            OpType::Expand,
            OpType::Tile,
            OpType::Pad,
            OpType::Slice,
            OpType::Split,
            OpType::SplitToSequence,
            OpType::Concat,
            OpType::ConcatFromSequence,
            OpType::DepthToSpace,
            OpType::SpaceToDepth,
            OpType::ReverseSequence,
            OpType::Gather,
            OpType::GatherElements,
            OpType::GatherND,
            OpType::Scatter,
            OpType::ScatterElements,
            OpType::ScatterND,
            OpType::ReduceMax,
            OpType::ReduceMin,
            OpType::ReduceMean,
            OpType::ReduceSum,
            OpType::ReduceProd,
            OpType::ReduceL1,
            OpType::ReduceL2,
            OpType::ReduceLogSum,
            OpType::ReduceLogSumExp,
            OpType::ReduceSumSquare,
            OpType::ArgMax,
            OpType::ArgMin,
            OpType::Equal,
            OpType::Greater,
            OpType::GreaterOrEqual,
            OpType::Less,
            OpType::LessOrEqual,
            OpType::Not,
            OpType::And,
            OpType::Or,
            OpType::Xor,
            OpType::Where,
            OpType::IsNaN,
            OpType::IsInf,
            OpType::Shape,
            OpType::Size,
            OpType::Cast,
            OpType::CastLike,
            OpType::ConstantOfShape,
            OpType::Range,
            OpType::OneHot,
            OpType::NonZero,
            OpType::TopK,
            OpType::Unique,
            OpType::EyeLike,
            OpType::Compress,
            OpType::LSTM,
            OpType::GRU,
            OpType::RNN,
            OpType::Constant,
            OpType::Identity,
            OpType::If,
            OpType::Loop,
            OpType::Scan,
            OpType::Resize,
            OpType::Upsample,
            OpType::Dropout,
            OpType::QuantizeLinear,
            OpType::DequantizeLinear,
            OpType::DynamicQuantizeLinear,
            OpType::QLinearConv,
            OpType::QLinearMatMul,
            OpType::Attention,
            OpType::Einsum,
            OpType::SequenceConstruct,
            OpType::SequenceAt,
            OpType::SequenceEmpty,
            OpType::SequenceInsert,
            OpType::SequenceErase,
            OpType::SequenceLength,
            OpType::SequenceMap,
            OpType::Optional,
            OpType::OptionalGetElement,
            OpType::OptionalHasElement,
            OpType::ImageDecoder,
            OpType::GridSample,
            OpType::RoiAlign,
            OpType::AffineGrid,
            OpType::CumSum,
            OpType::Det,
            OpType::Erf,
            OpType::Hardmax,
            OpType::Shrink,
            OpType::StringNormalizer,
            OpType::TfIdfVectorizer,
            OpType::NegativeLogLikelihoodLoss,
            OpType::SoftmaxCrossEntropyLoss,
            OpType::Bernoulli,
            OpType::RandomNormal,
            OpType::RandomNormalLike,
            OpType::RandomUniform,
            OpType::RandomUniformLike,
            OpType::Multinomial,
            OpType::CenterCropPad,
            OpType::Col2Im,
            OpType::Trilu,
            OpType::BlackmanWindow,
            OpType::HannWindow,
            OpType::HammingWindow,
            OpType::DFT,
            OpType::MelWeightMatrix,
            OpType::STFT,
            OpType::RegexFullMatch,
            OpType::StringConcat,
            OpType::StringSplit,
        ]
    }

    pub fn as_str(&self) -> &str {
        match self {
            OpType::Custom(s) => s,
            other => op_type_to_str(other),
        }
    }
}

impl Default for OpType<'_> {
    fn default() -> Self {
        OpType::Custom("")
    }
}

impl std::fmt::Display for OpType<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            OpType::Custom(s) => write!(f, "{s}"),
            other => write!(f, "{}", op_type_to_str(other)),
        }
    }
}

impl<'a> From<&'a str> for OpType<'a> {
    fn from(s: &'a str) -> Self {
        match s {
            "Relu" => OpType::Relu,
            "LeakyRelu" => OpType::LeakyRelu,
            "Sigmoid" => OpType::Sigmoid,
            "Tanh" => OpType::Tanh,
            "Softmax" => OpType::Softmax,
            "LogSoftmax" => OpType::LogSoftmax,
            "Elu" => OpType::Elu,
            "Selu" => OpType::Selu,
            "PRelu" => OpType::PRelu,
            "Gelu" => OpType::Gelu,
            "HardSigmoid" => OpType::HardSigmoid,
            "HardSwish" => OpType::HardSwish,
            "Softplus" => OpType::Softplus,
            "Softsign" => OpType::Softsign,
            "ThresholdedRelu" => OpType::ThresholdedRelu,
            "Celu" => OpType::Celu,
            "Mish" => OpType::Mish,
            "Add" => OpType::Add,
            "Sub" => OpType::Sub,
            "Mul" => OpType::Mul,
            "Div" => OpType::Div,
            "Neg" => OpType::Neg,
            "Abs" => OpType::Abs,
            "Sqrt" => OpType::Sqrt,
            "Exp" => OpType::Exp,
            "Log" => OpType::Log,
            "Pow" => OpType::Pow,
            "Ceil" => OpType::Ceil,
            "Floor" => OpType::Floor,
            "Round" => OpType::Round,
            "Clip" => OpType::Clip,
            "Sign" => OpType::Sign,
            "Reciprocal" => OpType::Reciprocal,
            "Mod" => OpType::Mod,
            "BitShift" => OpType::BitShift,
            "BitwiseAnd" => OpType::BitwiseAnd,
            "BitwiseOr" => OpType::BitwiseOr,
            "BitwiseXor" => OpType::BitwiseXor,
            "BitwiseNot" => OpType::BitwiseNot,
            "Min" => OpType::Min,
            "Max" => OpType::Max,
            "Mean" => OpType::Mean,
            "Sum" => OpType::Sum,
            "Sin" => OpType::Sin,
            "Cos" => OpType::Cos,
            "Tan" => OpType::Tan,
            "Asin" => OpType::Asin,
            "Acos" => OpType::Acos,
            "Atan" => OpType::Atan,
            "Sinh" => OpType::Sinh,
            "Cosh" => OpType::Cosh,
            "Asinh" => OpType::Asinh,
            "Acosh" => OpType::Acosh,
            "Atanh" => OpType::Atanh,
            "Conv" => OpType::Conv,
            "ConvTranspose" => OpType::ConvTranspose,
            "ConvInteger" => OpType::ConvInteger,
            "AveragePool" => OpType::AveragePool,
            "MaxPool" => OpType::MaxPool,
            "GlobalAveragePool" => OpType::GlobalAveragePool,
            "GlobalMaxPool" => OpType::GlobalMaxPool,
            "GlobalLpPool" => OpType::GlobalLpPool,
            "LpPool" => OpType::LpPool,
            "MaxRoiPool" => OpType::MaxRoiPool,
            "MaxUnpool" => OpType::MaxUnpool,
            "BatchNormalization" => OpType::BatchNormalization,
            "InstanceNormalization" => OpType::InstanceNormalization,
            "LayerNormalization" => OpType::LayerNormalization,
            "GroupNormalization" => OpType::GroupNormalization,
            "LpNormalization" => OpType::LpNormalization,
            "MeanVarianceNormalization" => OpType::MeanVarianceNormalization,
            "MatMul" => OpType::MatMul,
            "MatMulInteger" => OpType::MatMulInteger,
            "Gemm" => OpType::Gemm,
            "Reshape" => OpType::Reshape,
            "Transpose" => OpType::Transpose,
            "Flatten" => OpType::Flatten,
            "Squeeze" => OpType::Squeeze,
            "Unsqueeze" => OpType::Unsqueeze,
            "Expand" => OpType::Expand,
            "Tile" => OpType::Tile,
            "Pad" => OpType::Pad,
            "Slice" => OpType::Slice,
            "Split" => OpType::Split,
            "SplitToSequence" => OpType::SplitToSequence,
            "Concat" => OpType::Concat,
            "ConcatFromSequence" => OpType::ConcatFromSequence,
            "DepthToSpace" => OpType::DepthToSpace,
            "SpaceToDepth" => OpType::SpaceToDepth,
            "ReverseSequence" => OpType::ReverseSequence,
            "Gather" => OpType::Gather,
            "GatherElements" => OpType::GatherElements,
            "GatherND" => OpType::GatherND,
            "Scatter" => OpType::Scatter,
            "ScatterElements" => OpType::ScatterElements,
            "ScatterND" => OpType::ScatterND,
            "ReduceMax" => OpType::ReduceMax,
            "ReduceMin" => OpType::ReduceMin,
            "ReduceMean" => OpType::ReduceMean,
            "ReduceSum" => OpType::ReduceSum,
            "ReduceProd" => OpType::ReduceProd,
            "ReduceL1" => OpType::ReduceL1,
            "ReduceL2" => OpType::ReduceL2,
            "ReduceLogSum" => OpType::ReduceLogSum,
            "ReduceLogSumExp" => OpType::ReduceLogSumExp,
            "ReduceSumSquare" => OpType::ReduceSumSquare,
            "ArgMax" => OpType::ArgMax,
            "ArgMin" => OpType::ArgMin,
            "Equal" => OpType::Equal,
            "Greater" => OpType::Greater,
            "GreaterOrEqual" => OpType::GreaterOrEqual,
            "Less" => OpType::Less,
            "LessOrEqual" => OpType::LessOrEqual,
            "Not" => OpType::Not,
            "And" => OpType::And,
            "Or" => OpType::Or,
            "Xor" => OpType::Xor,
            "Where" => OpType::Where,
            "IsNaN" => OpType::IsNaN,
            "IsInf" => OpType::IsInf,
            "Shape" => OpType::Shape,
            "Size" => OpType::Size,
            "Cast" => OpType::Cast,
            "CastLike" => OpType::CastLike,
            "ConstantOfShape" => OpType::ConstantOfShape,
            "Range" => OpType::Range,
            "OneHot" => OpType::OneHot,
            "NonZero" => OpType::NonZero,
            "TopK" => OpType::TopK,
            "Unique" => OpType::Unique,
            "EyeLike" => OpType::EyeLike,
            "Compress" => OpType::Compress,
            "LSTM" => OpType::LSTM,
            "GRU" => OpType::GRU,
            "RNN" => OpType::RNN,
            "Constant" => OpType::Constant,
            "Identity" => OpType::Identity,
            "If" => OpType::If,
            "Loop" => OpType::Loop,
            "Scan" => OpType::Scan,
            "Resize" => OpType::Resize,
            "Upsample" => OpType::Upsample,
            "Dropout" => OpType::Dropout,
            "QuantizeLinear" => OpType::QuantizeLinear,
            "DequantizeLinear" => OpType::DequantizeLinear,
            "DynamicQuantizeLinear" => OpType::DynamicQuantizeLinear,
            "QLinearConv" => OpType::QLinearConv,
            "QLinearMatMul" => OpType::QLinearMatMul,
            "Attention" => OpType::Attention,
            "Einsum" => OpType::Einsum,
            "SequenceConstruct" => OpType::SequenceConstruct,
            "SequenceAt" => OpType::SequenceAt,
            "SequenceEmpty" => OpType::SequenceEmpty,
            "SequenceInsert" => OpType::SequenceInsert,
            "SequenceErase" => OpType::SequenceErase,
            "SequenceLength" => OpType::SequenceLength,
            "SequenceMap" => OpType::SequenceMap,
            "Optional" => OpType::Optional,
            "OptionalGetElement" => OpType::OptionalGetElement,
            "OptionalHasElement" => OpType::OptionalHasElement,
            "ImageDecoder" => OpType::ImageDecoder,
            "GridSample" => OpType::GridSample,
            "RoiAlign" => OpType::RoiAlign,
            "AffineGrid" => OpType::AffineGrid,
            "Cumsum" | "CumSum" => OpType::CumSum,
            "Det" => OpType::Det,
            "Erf" => OpType::Erf,
            "Hardmax" | "HardMax" => OpType::Hardmax,
            "Shrink" => OpType::Shrink,
            "StringNormalizer" => OpType::StringNormalizer,
            "TfIdfVectorizer" => OpType::TfIdfVectorizer,
            "NegativeLogLikelihoodLoss" => OpType::NegativeLogLikelihoodLoss,
            "SoftmaxCrossEntropyLoss" => OpType::SoftmaxCrossEntropyLoss,
            "Bernoulli" => OpType::Bernoulli,
            "RandomNormal" => OpType::RandomNormal,
            "RandomNormalLike" => OpType::RandomNormalLike,
            "RandomUniform" => OpType::RandomUniform,
            "RandomUniformLike" => OpType::RandomUniformLike,
            "Multinomial" => OpType::Multinomial,
            "CenterCropPad" => OpType::CenterCropPad,
            "Col2Im" => OpType::Col2Im,
            "Trilu" => OpType::Trilu,
            "BlackmanWindow" => OpType::BlackmanWindow,
            "HannWindow" => OpType::HannWindow,
            "HammingWindow" => OpType::HammingWindow,
            "DFT" => OpType::DFT,
            "MelWeightMatrix" => OpType::MelWeightMatrix,
            "STFT" => OpType::STFT,
            "RegexFullMatch" => OpType::RegexFullMatch,
            "StringConcat" => OpType::StringConcat,
            "StringSplit" => OpType::StringSplit,
            other => OpType::Custom(other),
        }
    }
}

fn op_type_to_str(op: &OpType<'_>) -> &'static str {
    match op {
        OpType::Relu => "Relu",
        OpType::LeakyRelu => "LeakyRelu",
        OpType::Sigmoid => "Sigmoid",
        OpType::Tanh => "Tanh",
        OpType::Softmax => "Softmax",
        OpType::LogSoftmax => "LogSoftmax",
        OpType::Elu => "Elu",
        OpType::Selu => "Selu",
        OpType::PRelu => "PRelu",
        OpType::Gelu => "Gelu",
        OpType::HardSigmoid => "HardSigmoid",
        OpType::HardSwish => "HardSwish",
        OpType::Softplus => "Softplus",
        OpType::Softsign => "Softsign",
        OpType::ThresholdedRelu => "ThresholdedRelu",
        OpType::Celu => "Celu",
        OpType::Mish => "Mish",
        OpType::Add => "Add",
        OpType::Sub => "Sub",
        OpType::Mul => "Mul",
        OpType::Div => "Div",
        OpType::Neg => "Neg",
        OpType::Abs => "Abs",
        OpType::Sqrt => "Sqrt",
        OpType::Exp => "Exp",
        OpType::Log => "Log",
        OpType::Pow => "Pow",
        OpType::Ceil => "Ceil",
        OpType::Floor => "Floor",
        OpType::Round => "Round",
        OpType::Clip => "Clip",
        OpType::Sign => "Sign",
        OpType::Reciprocal => "Reciprocal",
        OpType::Mod => "Mod",
        OpType::BitShift => "BitShift",
        OpType::BitwiseAnd => "BitwiseAnd",
        OpType::BitwiseOr => "BitwiseOr",
        OpType::BitwiseXor => "BitwiseXor",
        OpType::BitwiseNot => "BitwiseNot",
        OpType::Min => "Min",
        OpType::Max => "Max",
        OpType::Mean => "Mean",
        OpType::Sum => "Sum",
        OpType::Sin => "Sin",
        OpType::Cos => "Cos",
        OpType::Tan => "Tan",
        OpType::Asin => "Asin",
        OpType::Acos => "Acos",
        OpType::Atan => "Atan",
        OpType::Sinh => "Sinh",
        OpType::Cosh => "Cosh",
        OpType::Asinh => "Asinh",
        OpType::Acosh => "Acosh",
        OpType::Atanh => "Atanh",
        OpType::Conv => "Conv",
        OpType::ConvTranspose => "ConvTranspose",
        OpType::ConvInteger => "ConvInteger",
        OpType::AveragePool => "AveragePool",
        OpType::MaxPool => "MaxPool",
        OpType::GlobalAveragePool => "GlobalAveragePool",
        OpType::GlobalMaxPool => "GlobalMaxPool",
        OpType::GlobalLpPool => "GlobalLpPool",
        OpType::LpPool => "LpPool",
        OpType::MaxRoiPool => "MaxRoiPool",
        OpType::MaxUnpool => "MaxUnpool",
        OpType::BatchNormalization => "BatchNormalization",
        OpType::InstanceNormalization => "InstanceNormalization",
        OpType::LayerNormalization => "LayerNormalization",
        OpType::GroupNormalization => "GroupNormalization",
        OpType::LpNormalization => "LpNormalization",
        OpType::MeanVarianceNormalization => "MeanVarianceNormalization",
        OpType::MatMul => "MatMul",
        OpType::MatMulInteger => "MatMulInteger",
        OpType::Gemm => "Gemm",
        OpType::Reshape => "Reshape",
        OpType::Transpose => "Transpose",
        OpType::Flatten => "Flatten",
        OpType::Squeeze => "Squeeze",
        OpType::Unsqueeze => "Unsqueeze",
        OpType::Expand => "Expand",
        OpType::Tile => "Tile",
        OpType::Pad => "Pad",
        OpType::Slice => "Slice",
        OpType::Split => "Split",
        OpType::SplitToSequence => "SplitToSequence",
        OpType::Concat => "Concat",
        OpType::ConcatFromSequence => "ConcatFromSequence",
        OpType::DepthToSpace => "DepthToSpace",
        OpType::SpaceToDepth => "SpaceToDepth",
        OpType::ReverseSequence => "ReverseSequence",
        OpType::Gather => "Gather",
        OpType::GatherElements => "GatherElements",
        OpType::GatherND => "GatherND",
        OpType::Scatter => "Scatter",
        OpType::ScatterElements => "ScatterElements",
        OpType::ScatterND => "ScatterND",
        OpType::ReduceMax => "ReduceMax",
        OpType::ReduceMin => "ReduceMin",
        OpType::ReduceMean => "ReduceMean",
        OpType::ReduceSum => "ReduceSum",
        OpType::ReduceProd => "ReduceProd",
        OpType::ReduceL1 => "ReduceL1",
        OpType::ReduceL2 => "ReduceL2",
        OpType::ReduceLogSum => "ReduceLogSum",
        OpType::ReduceLogSumExp => "ReduceLogSumExp",
        OpType::ReduceSumSquare => "ReduceSumSquare",
        OpType::ArgMax => "ArgMax",
        OpType::ArgMin => "ArgMin",
        OpType::Equal => "Equal",
        OpType::Greater => "Greater",
        OpType::GreaterOrEqual => "GreaterOrEqual",
        OpType::Less => "Less",
        OpType::LessOrEqual => "LessOrEqual",
        OpType::Not => "Not",
        OpType::And => "And",
        OpType::Or => "Or",
        OpType::Xor => "Xor",
        OpType::Where => "Where",
        OpType::IsNaN => "IsNaN",
        OpType::IsInf => "IsInf",
        OpType::Shape => "Shape",
        OpType::Size => "Size",
        OpType::Cast => "Cast",
        OpType::CastLike => "CastLike",
        OpType::ConstantOfShape => "ConstantOfShape",
        OpType::Range => "Range",
        OpType::OneHot => "OneHot",
        OpType::NonZero => "NonZero",
        OpType::TopK => "TopK",
        OpType::Unique => "Unique",
        OpType::EyeLike => "EyeLike",
        OpType::Compress => "Compress",
        OpType::LSTM => "LSTM",
        OpType::GRU => "GRU",
        OpType::RNN => "RNN",
        OpType::Constant => "Constant",
        OpType::Identity => "Identity",
        OpType::If => "If",
        OpType::Loop => "Loop",
        OpType::Scan => "Scan",
        OpType::Resize => "Resize",
        OpType::Upsample => "Upsample",
        OpType::Dropout => "Dropout",
        OpType::QuantizeLinear => "QuantizeLinear",
        OpType::DequantizeLinear => "DequantizeLinear",
        OpType::DynamicQuantizeLinear => "DynamicQuantizeLinear",
        OpType::QLinearConv => "QLinearConv",
        OpType::QLinearMatMul => "QLinearMatMul",
        OpType::Attention => "Attention",
        OpType::Einsum => "Einsum",
        OpType::SequenceConstruct => "SequenceConstruct",
        OpType::SequenceAt => "SequenceAt",
        OpType::SequenceEmpty => "SequenceEmpty",
        OpType::SequenceInsert => "SequenceInsert",
        OpType::SequenceErase => "SequenceErase",
        OpType::SequenceLength => "SequenceLength",
        OpType::SequenceMap => "SequenceMap",
        OpType::Optional => "Optional",
        OpType::OptionalGetElement => "OptionalGetElement",
        OpType::OptionalHasElement => "OptionalHasElement",
        OpType::ImageDecoder => "ImageDecoder",
        OpType::GridSample => "GridSample",
        OpType::RoiAlign => "RoiAlign",
        OpType::AffineGrid => "AffineGrid",
        OpType::CumSum => "CumSum",
        OpType::Det => "Det",
        OpType::Erf => "Erf",
        OpType::Hardmax => "Hardmax",
        OpType::Shrink => "Shrink",
        OpType::StringNormalizer => "StringNormalizer",
        OpType::TfIdfVectorizer => "TfIdfVectorizer",
        OpType::NegativeLogLikelihoodLoss => "NegativeLogLikelihoodLoss",
        OpType::SoftmaxCrossEntropyLoss => "SoftmaxCrossEntropyLoss",
        OpType::Bernoulli => "Bernoulli",
        OpType::RandomNormal => "RandomNormal",
        OpType::RandomNormalLike => "RandomNormalLike",
        OpType::RandomUniform => "RandomUniform",
        OpType::RandomUniformLike => "RandomUniformLike",
        OpType::Multinomial => "Multinomial",
        OpType::CenterCropPad => "CenterCropPad",
        OpType::Col2Im => "Col2Im",
        OpType::Trilu => "Trilu",
        OpType::BlackmanWindow => "BlackmanWindow",
        OpType::HannWindow => "HannWindow",
        OpType::HammingWindow => "HammingWindow",
        OpType::DFT => "DFT",
        OpType::MelWeightMatrix => "MelWeightMatrix",
        OpType::STFT => "STFT",
        OpType::RegexFullMatch => "RegexFullMatch",
        OpType::StringConcat => "StringConcat",
        OpType::StringSplit => "StringSplit",
        OpType::Custom(_) => unreachable!(),
    }
}

/// Element data type for tensors.
///
/// Corresponds to `TensorProto.DataType` in the ONNX protobuf schema.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::DataType;
///
/// let dt = DataType::Float;
/// assert_ne!(dt, DataType::Int32);
///
/// // Convert from protobuf integer values
/// let dt = DataType::try_from(1).unwrap();
/// assert_eq!(dt, DataType::Float);
///
/// // Invalid values produce an error
/// assert!(DataType::try_from(999).is_err());
/// ```
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DataType {
    #[default]
    Undefined,
    Float,
    Uint8,
    Int8,
    Uint16,
    Int16,
    Int32,
    Int64,
    String,
    Bool,
    Float16,
    Double,
    Uint32,
    Uint64,
    Complex64,
    Complex128,
    Bfloat16,
    Float8e4m3fn,
    Float8e4m3fnuz,
    Float8e5m2,
    Float8e5m2fnuz,
    Uint4,
    Int4,
    Float4e2m1,
}

impl TryFrom<i32> for DataType {
    type Error = crate::error::Error;

    fn try_from(value: i32) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(DataType::Undefined),
            1 => Ok(DataType::Float),
            2 => Ok(DataType::Uint8),
            3 => Ok(DataType::Int8),
            4 => Ok(DataType::Uint16),
            5 => Ok(DataType::Int16),
            6 => Ok(DataType::Int32),
            7 => Ok(DataType::Int64),
            8 => Ok(DataType::String),
            9 => Ok(DataType::Bool),
            10 => Ok(DataType::Float16),
            11 => Ok(DataType::Double),
            12 => Ok(DataType::Uint32),
            13 => Ok(DataType::Uint64),
            14 => Ok(DataType::Complex64),
            15 => Ok(DataType::Complex128),
            16 => Ok(DataType::Bfloat16),
            17 => Ok(DataType::Float8e4m3fn),
            18 => Ok(DataType::Float8e4m3fnuz),
            19 => Ok(DataType::Float8e5m2),
            20 => Ok(DataType::Float8e5m2fnuz),
            21 => Ok(DataType::Uint4),
            22 => Ok(DataType::Int4),
            23 => Ok(DataType::Float4e2m1),
            _ => Err(crate::error::Error::InvalidEnumValue {
                enum_name: "DataType",
                value,
            }),
        }
    }
}

/// Discriminant for the kind of value stored in an [`Attribute`].
///
/// Corresponds to `AttributeProto.AttributeType` in the ONNX protobuf schema.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AttributeType {
    #[default]
    Undefined,
    Float,
    Int,
    String,
    Tensor,
    Graph,
    Floats,
    Ints,
    Strings,
    Tensors,
    Graphs,
    SparseTensor,
    SparseTensors,
    TypeProto,
    TypeProtos,
}

impl TryFrom<i32> for AttributeType {
    type Error = crate::error::Error;

    fn try_from(value: i32) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(AttributeType::Undefined),
            1 => Ok(AttributeType::Float),
            2 => Ok(AttributeType::Int),
            3 => Ok(AttributeType::String),
            4 => Ok(AttributeType::Tensor),
            5 => Ok(AttributeType::Graph),
            6 => Ok(AttributeType::Floats),
            7 => Ok(AttributeType::Ints),
            8 => Ok(AttributeType::Strings),
            9 => Ok(AttributeType::Tensors),
            10 => Ok(AttributeType::Graphs),
            11 => Ok(AttributeType::SparseTensor),
            12 => Ok(AttributeType::SparseTensors),
            13 => Ok(AttributeType::TypeProto),
            14 => Ok(AttributeType::TypeProtos),
            _ => Err(crate::error::Error::InvalidEnumValue {
                enum_name: "AttributeType",
                value,
            }),
        }
    }
}

/// Where tensor data is stored.
///
/// `Default` means inline in the protobuf; `External` means in a separate file
/// referenced by [`TensorProto::external_data`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DataLocation {
    #[default]
    Default,
    External,
}

impl TryFrom<i32> for DataLocation {
    type Error = crate::error::Error;

    fn try_from(value: i32) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(DataLocation::Default),
            1 => Ok(DataLocation::External),
            _ => Err(crate::error::Error::InvalidEnumValue {
                enum_name: "DataLocation",
                value,
            }),
        }
    }
}

/// A single dimension in a tensor shape — either a fixed value or a symbolic parameter.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::Dimension;
///
/// let fixed = Dimension::Value(224);
/// let dynamic = Dimension::Param("batch_size");
/// ```
#[derive(Debug, Clone, PartialEq)]
pub enum Dimension<'a> {
    Value(i64),
    Param(&'a str),
}

impl Default for Dimension<'_> {
    fn default() -> Self {
        Dimension::Value(0)
    }
}

/// A single dimension in a [`TensorShape`], combining a [`Dimension`] value with
/// an optional denotation string.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TensorShapeDimension<'a> {
    pub value: Dimension<'a>,
    pub denotation: &'a str,
}

/// The shape of a tensor, as a list of dimensions.
///
/// Corresponds to `TensorShapeProto` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TensorShape<'a> {
    pub dim: Vec<TensorShapeDimension<'a>>,
}

/// Describes a tensor type with an element data type and optional shape.
///
/// Corresponds to `TypeProto.Tensor` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TensorTypeProto<'a> {
    pub elem_type: DataType,
    pub shape: Option<TensorShape<'a>>,
}

/// Describes a sequence type whose elements have a given [`TypeProto`].
#[derive(Debug, Clone, PartialEq)]
pub struct SequenceTypeProto<'a> {
    pub elem_type: Box<TypeProto<'a>>,
}

/// Describes a map type with a key [`DataType`] and a value [`TypeProto`].
#[derive(Debug, Clone, PartialEq)]
pub struct MapTypeProto<'a> {
    pub key_type: DataType,
    pub value_type: Box<TypeProto<'a>>,
}

/// Describes an optional type whose element has a given [`TypeProto`].
#[derive(Debug, Clone, PartialEq)]
pub struct OptionalTypeProto<'a> {
    pub elem_type: Box<TypeProto<'a>>,
}

/// Describes a sparse tensor type with an element data type and optional shape.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct SparseTensorTypeProto<'a> {
    pub elem_type: DataType,
    pub shape: Option<TensorShape<'a>>,
}

/// The concrete type stored inside a [`TypeProto`].
#[derive(Debug, Clone, PartialEq)]
pub enum TypeValue<'a> {
    Tensor(TensorTypeProto<'a>),
    Sequence(SequenceTypeProto<'a>),
    Map(MapTypeProto<'a>),
    Optional(OptionalTypeProto<'a>),
    SparseTensor(SparseTensorTypeProto<'a>),
}

/// Describes the type of a graph input, output, or intermediate value.
///
/// Corresponds to `TypeProto` in the ONNX protobuf schema. The [`value`](TypeProto::value)
/// field holds the specific type variant (tensor, sequence, map, etc.).
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TypeProto<'a> {
    pub value: Option<TypeValue<'a>>,
    pub denotation: &'a str,
}

/// Type and shape information for a graph input, output, or intermediate value.
///
/// Corresponds to `ValueInfoProto` in the ONNX protobuf schema.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::*;
///
/// let vi = ValueInfo {
///     name: "input_0",
///     r#type: Some(TypeProto {
///         value: Some(TypeValue::Tensor(TensorTypeProto {
///             elem_type: DataType::Float,
///             shape: Some(TensorShape {
///                 dim: vec![
///                     TensorShapeDimension {
///                         value: Dimension::Param("batch"),
///                         denotation: "",
///                     },
///                     TensorShapeDimension {
///                         value: Dimension::Value(784),
///                         denotation: "",
///                     },
///                 ],
///             }),
///         })),
///         denotation: "",
///     }),
///     ..Default::default()
/// };
///
/// assert_eq!(vi.name, "input_0");
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ValueInfo<'a> {
    pub name: &'a str,
    pub r#type: Option<TypeProto<'a>>,
    pub doc_string: &'a str,
    pub metadata_props: Vec<StringStringEntry<'a>>,
}

/// Byte range within a segmented tensor. Used for very large tensors that
/// are split across multiple `TensorProto` messages.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TensorSegment {
    pub begin: i64,
    pub end: i64,
}

/// An ONNX tensor value.
///
/// Corresponds to `TensorProto` in the ONNX protobuf schema. Tensor data
/// can be stored in typed fields or as raw bytes. Use the constructor methods
/// to create tensors and the accessor methods ([`as_f32`](TensorProto::as_f32),
/// [`as_i64`](TensorProto::as_i64), etc.) to read data.
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::*;
///
/// let tensor = TensorProto::from_f32("weights", vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
///
/// assert_eq!(tensor.dims(), &[2, 3]);
/// assert_eq!(tensor.as_f32().unwrap().len(), 6);
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TensorProto<'a> {
    pub(crate) dims: Vec<i64>,
    pub(crate) data_type: DataType,
    pub(crate) segment: Option<TensorSegment>,
    pub(crate) float_data: Vec<f32>,
    pub(crate) int32_data: Vec<i32>,
    pub(crate) string_data: Vec<&'a [u8]>,
    pub(crate) int64_data: Vec<i64>,
    pub(crate) raw_data: &'a [u8],
    pub(crate) double_data: Vec<f64>,
    pub(crate) uint64_data: Vec<u64>,
    pub(crate) name: &'a str,
    pub(crate) doc_string: &'a str,
    pub(crate) external_data: Vec<StringStringEntry<'a>>,
    pub(crate) data_location: DataLocation,
    pub(crate) metadata_props: Vec<StringStringEntry<'a>>,
}

/// A sparse tensor stored as separate values and indices tensors.
///
/// Corresponds to `SparseTensorProto` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct SparseTensor<'a> {
    pub values: Option<TensorProto<'a>>,
    pub indices: Option<TensorProto<'a>>,
    pub dims: Vec<i64>,
}

/// Quantization annotation for a tensor in the graph.
///
/// Corresponds to `TensorAnnotation` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TensorAnnotation<'a> {
    pub tensor_name: &'a str,
    pub quant_parameter_tensor_names: Vec<StringStringEntry<'a>>,
}

/// A named attribute on a graph [`Node`].
///
/// Corresponds to `AttributeProto` in the ONNX protobuf schema. The value is
/// stored in whichever field matches the [`type`](Attribute::type) discriminant
/// (e.g., `f` for floats, `i` for ints, `s` for byte strings).
///
/// # Examples
///
/// ```
/// use onnx_rs::ast::*;
///
/// let attr = Attribute {
///     name: "axis",
///     r#type: AttributeType::Int,
///     i: 1,
///     ..Default::default()
/// };
///
/// assert_eq!(attr.name, "axis");
/// assert_eq!(attr.i, 1);
/// ```
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Attribute<'a> {
    pub name: &'a str,
    pub ref_attr_name: &'a str,
    pub doc_string: &'a str,
    pub r#type: AttributeType,
    pub f: f32,
    pub i: i64,
    pub s: &'a [u8],
    pub t: Option<TensorProto<'a>>,
    pub g: Option<Box<Graph<'a>>>,
    pub sparse_tensor: Option<SparseTensor<'a>>,
    pub tp: Option<TypeProto<'a>>,
    pub floats: Vec<f32>,
    pub ints: Vec<i64>,
    pub strings: Vec<&'a [u8]>,
    pub tensors: Vec<TensorProto<'a>>,
    pub graphs: Vec<Graph<'a>>,
    pub sparse_tensors: Vec<SparseTensor<'a>>,
    pub type_protos: Vec<TypeProto<'a>>,
}

/// Training-specific information attached to a model.
///
/// Corresponds to `TrainingInfoProto` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TrainingInfo<'a> {
    pub initialization: Option<Graph<'a>>,
    pub algorithm: Option<Graph<'a>>,
    pub initialization_binding: Vec<StringStringEntry<'a>>,
    pub update_binding: Vec<StringStringEntry<'a>>,
}

/// A reusable function definition within an ONNX model.
///
/// Corresponds to `FunctionProto` in the ONNX protobuf schema.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Function<'a> {
    pub name: &'a str,
    pub input: Vec<&'a str>,
    pub output: Vec<&'a str>,
    pub attribute: Vec<&'a str>,
    pub attribute_proto: Vec<Attribute<'a>>,
    pub node: Vec<Node<'a>>,
    pub doc_string: &'a str,
    pub opset_import: Vec<OperatorSetId<'a>>,
    pub domain: &'a str,
    pub overload: &'a str,
    pub value_info: Vec<ValueInfo<'a>>,
    pub metadata_props: Vec<StringStringEntry<'a>>,
}

// ---------------------------------------------------------------------------
// Convenience methods
// ---------------------------------------------------------------------------

impl<'a> TensorProto<'a> {
    // -- Constructors --

    /// Creates a tensor with [`DataType::Float`] data.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_f32("W", vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    /// assert_eq!(t.name(), "W");
    /// assert_eq!(t.dims(), &[2, 3]);
    /// assert_eq!(t.data_type(), DataType::Float);
    /// assert_eq!(t.as_f32().unwrap().len(), 6);
    /// ```
    pub fn from_f32(name: &'a str, dims: Vec<i64>, data: Vec<f32>) -> Self {
        Self { name, dims, data_type: DataType::Float, float_data: data, ..Default::default() }
    }

    /// Creates a tensor with [`DataType::Double`] data.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_f64("bias", vec![2], vec![0.1, 0.2]);
    /// assert_eq!(t.data_type(), DataType::Double);
    /// assert_eq!(*t.as_f64().unwrap(), [0.1, 0.2]);
    /// ```
    pub fn from_f64(name: &'a str, dims: Vec<i64>, data: Vec<f64>) -> Self {
        Self { name, dims, data_type: DataType::Double, double_data: data, ..Default::default() }
    }

    /// Creates a tensor with [`DataType::Int32`] data.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_i32("idx", vec![3], vec![1, 2, 3]);
    /// assert_eq!(t.data_type(), DataType::Int32);
    /// assert_eq!(*t.as_i32().unwrap(), [1, 2, 3]);
    /// ```
    pub fn from_i32(name: &'a str, dims: Vec<i64>, data: Vec<i32>) -> Self {
        Self { name, dims, data_type: DataType::Int32, int32_data: data, ..Default::default() }
    }

    /// Creates a tensor with [`DataType::Int64`] data.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_i64("shape", vec![3], vec![1, 224, 224]);
    /// assert_eq!(t.data_type(), DataType::Int64);
    /// assert_eq!(*t.as_i64().unwrap(), [1, 224, 224]);
    /// ```
    pub fn from_i64(name: &'a str, dims: Vec<i64>, data: Vec<i64>) -> Self {
        Self { name, dims, data_type: DataType::Int64, int64_data: data, ..Default::default() }
    }

    /// Creates a tensor with [`DataType::Uint64`] data.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_u64("counts", vec![2], vec![42, 100]);
    /// assert_eq!(t.data_type(), DataType::Uint64);
    /// assert_eq!(*t.as_u64().unwrap(), [42, 100]);
    /// ```
    pub fn from_u64(name: &'a str, dims: Vec<i64>, data: Vec<u64>) -> Self {
        Self { name, dims, data_type: DataType::Uint64, uint64_data: data, ..Default::default() }
    }

    /// Creates a tensor with [`DataType::String`] data.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_strings("labels", vec![2], vec![b"cat", b"dog"]);
    /// assert_eq!(t.data_type(), DataType::String);
    /// assert_eq!(t.as_strings().unwrap(), &[b"cat".as_slice(), b"dog".as_slice()]);
    /// ```
    pub fn from_strings(name: &'a str, dims: Vec<i64>, data: Vec<&'a [u8]>) -> Self {
        Self { name, dims, data_type: DataType::String, string_data: data, ..Default::default() }
    }

    /// Creates a tensor from raw bytes with a given [`DataType`].
    ///
    /// Use this for types stored as raw little-endian bytes (the default for
    /// modern exporters like PyTorch), or for exotic types like Float8 or Int4.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let bytes = [1.0f32, 2.0].iter().flat_map(|f| f.to_le_bytes()).collect::<Vec<_>>();
    /// let t = TensorProto::from_raw("W", vec![2], DataType::Float, &bytes);
    /// assert_eq!(*t.as_f32().unwrap(), [1.0, 2.0]);
    /// ```
    pub fn from_raw(name: &'a str, dims: Vec<i64>, data_type: DataType, data: &'a [u8]) -> Self {
        Self { name, dims, data_type, raw_data: data, ..Default::default() }
    }

    // -- Getters --

    /// Returns the tensor name.
    ///
    /// ```
    /// use onnx_rs::ast::TensorProto;
    ///
    /// let t = TensorProto::from_f32("weight", vec![2], vec![1.0, 2.0]);
    /// assert_eq!(t.name(), "weight");
    /// ```
    pub fn name(&self) -> &'a str {
        self.name
    }

    /// Returns the tensor dimensions.
    ///
    /// ```
    /// use onnx_rs::ast::TensorProto;
    ///
    /// let t = TensorProto::from_f32("W", vec![2, 3], vec![0.0; 6]);
    /// assert_eq!(t.dims(), &[2, 3]);
    /// ```
    pub fn dims(&self) -> &[i64] {
        &self.dims
    }

    /// Returns the tensor's [`DataType`].
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_i64("ids", vec![3], vec![1, 2, 3]);
    /// assert_eq!(t.data_type(), DataType::Int64);
    /// ```
    pub fn data_type(&self) -> DataType {
        self.data_type
    }

    /// Returns the optional [`TensorSegment`] for segmented tensors.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("chunk", vec![], DataType::Float, &[0u8; 100])
    ///     .with_segment(TensorSegment { begin: 0, end: 50 });
    /// assert_eq!(t.segment().unwrap().begin, 0);
    /// ```
    pub fn segment(&self) -> Option<&TensorSegment> {
        self.segment.as_ref()
    }

    /// Returns the tensor's documentation string.
    ///
    /// ```
    /// use onnx_rs::ast::TensorProto;
    ///
    /// let t = TensorProto::from_f32("W", vec![], vec![]).with_doc_string("layer 1 weights");
    /// assert_eq!(t.doc_string(), "layer 1 weights");
    /// ```
    pub fn doc_string(&self) -> &'a str {
        self.doc_string
    }

    /// Returns the external data entries for externally stored tensors.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("ext", vec![], DataType::Float, &[])
    ///     .with_data_location(DataLocation::External)
    ///     .with_external_data(vec![StringStringEntry { key: "location", value: "w.bin" }]);
    /// assert_eq!(t.external_data()[0].key, "location");
    /// ```
    pub fn external_data(&self) -> &[StringStringEntry<'a>] {
        &self.external_data
    }

    /// Returns where the tensor data is stored ([`DataLocation::Default`] or
    /// [`DataLocation::External`]).
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("ext", vec![], DataType::Float, &[])
    ///     .with_data_location(DataLocation::External);
    /// assert_eq!(t.data_location(), DataLocation::External);
    /// ```
    pub fn data_location(&self) -> DataLocation {
        self.data_location
    }

    /// Returns the metadata key-value pairs attached to this tensor.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_f32("W", vec![], vec![])
    ///     .with_metadata_props(vec![StringStringEntry { key: "author", value: "onnx-rs" }]);
    /// assert_eq!(t.metadata_props()[0].value, "onnx-rs");
    /// ```
    pub fn metadata_props(&self) -> &[StringStringEntry<'a>] {
        &self.metadata_props
    }

    // -- Builder-style setters --

    /// Sets the [`TensorSegment`] for segmented tensors.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("chunk", vec![], DataType::Float, &[0u8; 100])
    ///     .with_segment(TensorSegment { begin: 0, end: 50 });
    /// assert_eq!(t.segment().unwrap().end, 50);
    /// ```
    pub fn with_segment(mut self, segment: TensorSegment) -> Self {
        self.segment = Some(segment);
        self
    }

    /// Sets the documentation string.
    ///
    /// ```
    /// use onnx_rs::ast::TensorProto;
    ///
    /// let t = TensorProto::from_f32("W", vec![], vec![]).with_doc_string("my doc");
    /// assert_eq!(t.doc_string(), "my doc");
    /// ```
    pub fn with_doc_string(mut self, doc_string: &'a str) -> Self {
        self.doc_string = doc_string;
        self
    }

    /// Sets external data entries for externally stored tensors.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("ext", vec![], DataType::Float, &[])
    ///     .with_external_data(vec![StringStringEntry { key: "location", value: "w.bin" }]);
    /// assert_eq!(t.external_data().len(), 1);
    /// ```
    pub fn with_external_data(mut self, external_data: Vec<StringStringEntry<'a>>) -> Self {
        self.external_data = external_data;
        self
    }

    /// Sets the data location ([`DataLocation::Default`] or [`DataLocation::External`]).
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("ext", vec![], DataType::Float, &[])
    ///     .with_data_location(DataLocation::External);
    /// assert_eq!(t.data_location(), DataLocation::External);
    /// ```
    pub fn with_data_location(mut self, data_location: DataLocation) -> Self {
        self.data_location = data_location;
        self
    }

    /// Sets metadata key-value pairs on the tensor.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_f32("W", vec![], vec![])
    ///     .with_metadata_props(vec![StringStringEntry { key: "k", value: "v" }]);
    /// assert_eq!(t.metadata_props().len(), 1);
    /// ```
    pub fn with_metadata_props(mut self, metadata_props: Vec<StringStringEntry<'a>>) -> Self {
        self.metadata_props = metadata_props;
        self
    }

    // -- Data accessors --

    /// Returns `f32` data, reading from either the typed field or raw bytes.
    ///
    /// Returns [`Cow::Borrowed`](std::borrow::Cow::Borrowed) when data is in the
    /// typed field (zero-copy), or [`Cow::Owned`](std::borrow::Cow::Owned) when
    /// decoded from raw bytes.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_f32("W", vec![2], vec![1.0, 2.0]);
    /// assert_eq!(*t.as_f32().unwrap(), [1.0, 2.0]);
    ///
    /// let bytes: Vec<u8> = [3.0f32, 4.0].iter().flat_map(|f| f.to_le_bytes()).collect();
    /// let t = TensorProto::from_raw("W", vec![2], DataType::Float, &bytes);
    /// assert_eq!(*t.as_f32().unwrap(), [3.0, 4.0]);
    /// ```
    pub fn as_f32(&self) -> Option<std::borrow::Cow<'_, [f32]>> {
        if !self.float_data.is_empty() {
            return Some(std::borrow::Cow::Borrowed(&self.float_data));
        }
        if !self.raw_data.is_empty() && self.data_type == DataType::Float {
            return Some(std::borrow::Cow::Owned(decode_raw_f32(self.raw_data)));
        }
        None
    }

    /// Returns `f64` data, reading from either the typed field or raw bytes.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_f64("b", vec![2], vec![0.1, 0.2]);
    /// assert_eq!(*t.as_f64().unwrap(), [0.1, 0.2]);
    /// ```
    pub fn as_f64(&self) -> Option<std::borrow::Cow<'_, [f64]>> {
        if !self.double_data.is_empty() {
            return Some(std::borrow::Cow::Borrowed(&self.double_data));
        }
        if !self.raw_data.is_empty() && self.data_type == DataType::Double {
            return Some(std::borrow::Cow::Owned(decode_raw_f64(self.raw_data)));
        }
        None
    }

    /// Returns `i32` data, reading from either the typed field or raw bytes.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_i32("x", vec![3], vec![1, -2, 3]);
    /// assert_eq!(*t.as_i32().unwrap(), [1, -2, 3]);
    /// ```
    pub fn as_i32(&self) -> Option<std::borrow::Cow<'_, [i32]>> {
        if !self.int32_data.is_empty() {
            return Some(std::borrow::Cow::Borrowed(&self.int32_data));
        }
        if !self.raw_data.is_empty() && self.data_type == DataType::Int32 {
            return Some(std::borrow::Cow::Owned(decode_raw_i32(self.raw_data)));
        }
        None
    }

    /// Returns `i64` data, reading from either the typed field or raw bytes.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_i64("ids", vec![3], vec![10, 20, 30]);
    /// assert_eq!(*t.as_i64().unwrap(), [10, 20, 30]);
    /// ```
    pub fn as_i64(&self) -> Option<std::borrow::Cow<'_, [i64]>> {
        if !self.int64_data.is_empty() {
            return Some(std::borrow::Cow::Borrowed(&self.int64_data));
        }
        if !self.raw_data.is_empty() && self.data_type == DataType::Int64 {
            return Some(std::borrow::Cow::Owned(decode_raw_i64(self.raw_data)));
        }
        None
    }

    /// Returns `u64` data, reading from either the typed field or raw bytes.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_u64("counts", vec![2], vec![42, 100]);
    /// assert_eq!(*t.as_u64().unwrap(), [42, 100]);
    /// ```
    pub fn as_u64(&self) -> Option<std::borrow::Cow<'_, [u64]>> {
        if !self.uint64_data.is_empty() {
            return Some(std::borrow::Cow::Borrowed(&self.uint64_data));
        }
        if !self.raw_data.is_empty() && self.data_type == DataType::Uint64 {
            return Some(std::borrow::Cow::Owned(decode_raw_u64(self.raw_data)));
        }
        None
    }

    /// Returns string data as byte slices.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_strings("labels", vec![2], vec![b"yes", b"no"]);
    /// assert_eq!(t.as_strings().unwrap().len(), 2);
    /// ```
    pub fn as_strings(&self) -> Option<&[&'a [u8]]> {
        if !self.string_data.is_empty() {
            return Some(&self.string_data);
        }
        None
    }

    /// Returns the raw byte data, if present.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let t = TensorProto::from_raw("x", vec![4], DataType::Float8e4m3fn, &[1, 2, 3, 4]);
    /// assert_eq!(t.as_raw().unwrap(), &[1, 2, 3, 4]);
    /// ```
    pub fn as_raw(&self) -> Option<&'a [u8]> {
        if !self.raw_data.is_empty() {
            return Some(self.raw_data);
        }
        None
    }
}

fn decode_raw_f32(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(4)
        .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
        .collect()
}

fn decode_raw_f64(bytes: &[u8]) -> Vec<f64> {
    bytes
        .chunks_exact(8)
        .map(|c| f64::from_le_bytes(c.try_into().unwrap()))
        .collect()
}

fn decode_raw_i32(bytes: &[u8]) -> Vec<i32> {
    bytes
        .chunks_exact(4)
        .map(|c| i32::from_le_bytes(c.try_into().unwrap()))
        .collect()
}

fn decode_raw_i64(bytes: &[u8]) -> Vec<i64> {
    bytes
        .chunks_exact(8)
        .map(|c| i64::from_le_bytes(c.try_into().unwrap()))
        .collect()
}

fn decode_raw_u64(bytes: &[u8]) -> Vec<u64> {
    bytes
        .chunks_exact(8)
        .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
        .collect()
}

impl<'a> Graph<'a> {
    /// Finds an initializer tensor by name.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let graph = Graph {
    ///     initializer: vec![
    ///         TensorProto::from_f32("W", vec![2], vec![1.0, 2.0]),
    ///         TensorProto::from_f32("b", vec![1], vec![0.5]),
    ///     ],
    ///     ..Default::default()
    /// };
    /// assert_eq!(graph.get_initializer("b").unwrap().dims(), &[1]);
    /// assert!(graph.get_initializer("missing").is_none());
    /// ```
    pub fn get_initializer(&self, name: &str) -> Option<&TensorProto<'a>> {
        self.initializer.iter().find(|t| t.name == name)
    }

    /// Returns graph inputs that are not also initializers (the true runtime inputs).
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let graph = Graph {
    ///     input: vec![
    ///         ValueInfo { name: "X", ..Default::default() },
    ///         ValueInfo { name: "W", ..Default::default() },
    ///     ],
    ///     initializer: vec![TensorProto::from_f32("W", vec![2], vec![1.0, 2.0])],
    ///     ..Default::default()
    /// };
    /// let inputs = graph.non_init_inputs();
    /// assert_eq!(inputs.len(), 1);
    /// assert_eq!(inputs[0].name, "X");
    /// ```
    pub fn non_init_inputs(&self) -> Vec<&ValueInfo<'a>> {
        self.input
            .iter()
            .filter(|vi| !self.initializer.iter().any(|init| init.name == vi.name))
            .collect()
    }
}

impl<'a> ValueInfo<'a> {
    /// Returns the element [`DataType`] if this value is a tensor type.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let vi = ValueInfo {
    ///     name: "X",
    ///     r#type: Some(TypeProto {
    ///         value: Some(TypeValue::Tensor(TensorTypeProto {
    ///             elem_type: DataType::Float,
    ///             shape: None,
    ///         })),
    ///         denotation: "",
    ///     }),
    ///     ..Default::default()
    /// };
    /// assert_eq!(vi.tensor_elem_type(), Some(DataType::Float));
    /// ```
    pub fn tensor_elem_type(&self) -> Option<DataType> {
        match &self.r#type {
            Some(tp) => match &tp.value {
                Some(TypeValue::Tensor(t)) => Some(t.elem_type),
                _ => None,
            },
            None => None,
        }
    }

    /// Returns the [`TensorShape`] if this value is a tensor type with a shape.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let vi = ValueInfo {
    ///     name: "X",
    ///     r#type: Some(TypeProto {
    ///         value: Some(TypeValue::Tensor(TensorTypeProto {
    ///             elem_type: DataType::Float,
    ///             shape: Some(TensorShape {
    ///                 dim: vec![
    ///                     TensorShapeDimension { value: Dimension::Param("N"), denotation: "" },
    ///                     TensorShapeDimension { value: Dimension::Value(784), denotation: "" },
    ///                 ],
    ///             }),
    ///         })),
    ///         denotation: "",
    ///     }),
    ///     ..Default::default()
    /// };
    /// assert_eq!(vi.tensor_shape().unwrap().rank(), 2);
    /// ```
    pub fn tensor_shape(&self) -> Option<&TensorShape<'a>> {
        match &self.r#type {
            Some(tp) => match &tp.value {
                Some(TypeValue::Tensor(t)) => t.shape.as_ref(),
                _ => None,
            },
            None => None,
        }
    }
}

impl TensorShape<'_> {
    /// Returns the number of dimensions (rank) of the tensor shape.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let shape = TensorShape {
    ///     dim: vec![
    ///         TensorShapeDimension { value: Dimension::Value(3), denotation: "" },
    ///         TensorShapeDimension { value: Dimension::Value(224), denotation: "" },
    ///     ],
    /// };
    /// assert_eq!(shape.rank(), 2);
    /// ```
    pub fn rank(&self) -> usize {
        self.dim.len()
    }

    /// Returns concrete dimension values if all dimensions are fixed.
    ///
    /// Returns `None` if any dimension is a symbolic parameter.
    ///
    /// ```
    /// use onnx_rs::ast::*;
    ///
    /// let fixed = TensorShape {
    ///     dim: vec![
    ///         TensorShapeDimension { value: Dimension::Value(3), denotation: "" },
    ///         TensorShapeDimension { value: Dimension::Value(224), denotation: "" },
    ///     ],
    /// };
    /// assert_eq!(fixed.fixed_dims(), Some(vec![3, 224]));
    ///
    /// let dynamic = TensorShape {
    ///     dim: vec![
    ///         TensorShapeDimension { value: Dimension::Param("batch"), denotation: "" },
    ///         TensorShapeDimension { value: Dimension::Value(224), denotation: "" },
    ///     ],
    /// };
    /// assert_eq!(dynamic.fixed_dims(), None);
    /// ```
    pub fn fixed_dims(&self) -> Option<Vec<i64>> {
        self.dim
            .iter()
            .map(|d| match &d.value {
                Dimension::Value(v) => Some(*v),
                Dimension::Param(_) => None,
            })
            .collect()
    }
}