// ONNX Protocol Buffer definitions
// Based on ONNX spec: https://github.com/onnx/onnx/blob/main/onnx/onnx.proto
// Simplified for RustyTorch code generation

syntax = "proto3";

package onnx;

// Tensor data types
enum TensorProto_DataType {
  UNDEFINED = 0;
  FLOAT = 1;
  UINT8 = 2;
  INT8 = 3;
  UINT16 = 4;
  INT16 = 5;
  INT32 = 6;
  INT64 = 7;
  STRING = 8;
  BOOL = 9;
  FLOAT16 = 10;
  DOUBLE = 11;
  UINT32 = 12;
  UINT64 = 13;
  COMPLEX64 = 14;
  COMPLEX128 = 15;
  BFLOAT16 = 16;
  FLOAT8E4M3FN = 17;
  FLOAT8E4M3FNUZ = 18;
  FLOAT8E5M2 = 19;
  FLOAT8E5M2FNUZ = 20;
  UINT4 = 21;
  INT4 = 22;
}

// Attribute types
enum AttributeProto_AttributeType {
  UNDEFINED_ATTR = 0;
  FLOAT_ATTR = 1;
  INT_ATTR = 2;
  STRING_ATTR = 3;
  TENSOR_ATTR = 4;
  GRAPH_ATTR = 5;
  SPARSE_TENSOR_ATTR = 11;
  TYPE_PROTO_ATTR = 13;
  FLOATS_ATTR = 6;
  INTS_ATTR = 7;
  STRINGS_ATTR = 8;
  TENSORS_ATTR = 9;
  GRAPHS_ATTR = 10;
  SPARSE_TENSORS_ATTR = 12;
  TYPE_PROTOS_ATTR = 14;
}

// Tensor shape dimension
message TensorShapeProto {
  message Dimension {
    oneof value {
      int64 dim_value = 1;
      string dim_param = 2;
    }
    string denotation = 3;
  }
  repeated Dimension dim = 1;
}

// Type information for tensors
message TypeProto {
  message Tensor {
    int32 elem_type = 1;
    TensorShapeProto shape = 2;
  }

  message Sequence {
    TypeProto elem_type = 1;
  }

  message Map {
    int32 key_type = 1;
    TypeProto value_type = 2;
  }

  message Optional {
    TypeProto elem_type = 1;
  }

  message SparseTensor {
    int32 elem_type = 1;
    TensorShapeProto shape = 2;
  }

  oneof value {
    Tensor tensor_type = 1;
    Sequence sequence_type = 4;
    Map map_type = 5;
    Optional optional_type = 9;
    SparseTensor sparse_tensor_type = 8;
  }

  string denotation = 6;
}

// Tensor data (weights, constants)
message TensorProto {
  repeated int64 dims = 1;
  int32 data_type = 2;

  message Segment {
    int64 begin = 1;
    int64 end = 2;
  }
  Segment segment = 3;

  // Raw tensor data (for all types)
  repeated float float_data = 4 [packed = true];
  repeated int32 int32_data = 5 [packed = true];
  repeated bytes string_data = 6;
  repeated int64 int64_data = 7 [packed = true];

  string name = 8;
  string doc_string = 12;

  bytes raw_data = 9;

  repeated StringStringEntryProto external_data = 13;
  int32 data_location = 14;

  repeated double double_data = 10 [packed = true];
  repeated uint64 uint64_data = 11 [packed = true];
}

// Sparse tensor
message SparseTensorProto {
  repeated int64 dims = 1;
  TensorProto indices = 2;
  TensorProto values = 3;
}

// Attribute (node parameters)
message AttributeProto {
  string name = 1;
  string ref_attr_name = 21;
  string doc_string = 13;
  int32 type = 20;

  // Single values
  float f = 2;
  int64 i = 3;
  bytes s = 4;
  TensorProto t = 5;
  GraphProto g = 6;
  SparseTensorProto sparse_tensor = 22;
  TypeProto tp = 14;

  // Repeated values
  repeated float floats = 7;
  repeated int64 ints = 8;
  repeated bytes strings = 9;
  repeated TensorProto tensors = 10;
  repeated GraphProto graphs = 11;
  repeated SparseTensorProto sparse_tensors = 23;
  repeated TypeProto type_protos = 15;
}

// Graph input/output
message ValueInfoProto {
  string name = 1;
  TypeProto type = 2;
  string doc_string = 3;
}

// Graph node (operation)
message NodeProto {
  repeated string input = 1;
  repeated string output = 2;
  string name = 3;
  string op_type = 4;
  string domain = 7;
  repeated AttributeProto attribute = 5;
  string doc_string = 6;
}

// Computation graph
message GraphProto {
  repeated NodeProto node = 1;
  string name = 2;
  repeated TensorProto initializer = 5;
  repeated SparseTensorProto sparse_initializer = 15;
  string doc_string = 10;
  repeated ValueInfoProto input = 11;
  repeated ValueInfoProto output = 12;
  repeated ValueInfoProto value_info = 13;
  repeated TensorAnnotation quantization_annotation = 14;
}

// String-string key-value pair
message StringStringEntryProto {
  string key = 1;
  string value = 2;
}

// Tensor annotation
message TensorAnnotation {
  string tensor_name = 1;
  repeated StringStringEntryProto quant_parameter_tensor_names = 2;
}

// Operator set import
message OperatorSetIdProto {
  string domain = 1;
  int64 version = 2;
}

// Training info (optional)
message TrainingInfoProto {
  GraphProto initialization = 1;
  GraphProto algorithm = 2;
  repeated StringStringEntryProto initialization_binding = 3;
  repeated StringStringEntryProto update_binding = 4;
}

// Function definition
message FunctionProto {
  string name = 1;
  repeated string input = 4;
  repeated string output = 5;
  repeated string attribute = 6;
  repeated AttributeProto attribute_proto = 11;
  repeated NodeProto node = 7;
  string doc_string = 8;
  repeated OperatorSetIdProto opset_import = 9;
  string domain = 10;
}

// Model (top-level container)
message ModelProto {
  int64 ir_version = 1;
  repeated OperatorSetIdProto opset_import = 8;
  string producer_name = 2;
  string producer_version = 3;
  string domain = 4;
  int64 model_version = 5;
  string doc_string = 6;
  GraphProto graph = 7;
  repeated StringStringEntryProto metadata_props = 14;
  repeated TrainingInfoProto training_info = 20;
  repeated FunctionProto functions = 25;
}