//
//  OnnxMLIO.h
//  SwitchboardOnnx
//
//  Created by Tayyab Javed on 2024-07-02.
//

#pragma once

#include <string>
#include <iostream>
#include <vector>
#include <sstream>

namespace switchboard::extensions::onnx {

typedef enum {
    DATA_TYPE_UNDEFINED = 0,
    DATA_TYPE_FLOAT = 1,
    DATA_TYPE_UINT8 = 2,
    DATA_TYPE_INT8 = 3,
    DATA_TYPE_UINT16 = 4,
    DATA_TYPE_INT16 = 5,
    DATA_TYPE_INT32 = 6,
    DATA_TYPE_INT64 = 7,
    DATA_TYPE_STRING = 8,
    DATA_TYPE_BOOL = 9,
    DATA_TYPE_FLOAT16 = 10,
    DATA_TYPE_DOUBLE = 11,
    DATA_TYPE_UINT32 = 12,
    DATA_TYPE_UINT64 = 13,
    DATA_TYPE_COMPLEX64 = 14,
    DATA_TYPE_COMPLEX128 = 15,
    DATA_TYPE_BFLOAT16 = 16
} DataType;



struct TensorInfo {
    std::string name;
    size_t rank;
    std::vector<int64_t> shape;
    DataType dataType;

    static std::string shapeToString(const std::vector<int64_t>& shape) {
        std::stringstream ss;
        ss << "[";
        for (size_t i = 0; i < shape.size(); ++i) {
            ss << shape[i];

            if (i < shape.size() - 1) {
                ss << ", ";
            }
        }
        ss << "]";
        return ss.str();
    }

    std::string toString() {
        return "name: " + name + ", " +
               "dataType = " + std::to_string(dataType) + ", " +
               "rank = " + std::to_string(rank) + ", " +
               "shapeToString = " + shapeToString(shape);
    }
};

}
