//
//  ModelInstance.hpp
//  SwitchboardOnnx
//
//  Created by Tayyab Javed on 2024-05-29.
//

#pragma once

#include "OnnxMLIO.h"
#include "OnnxProviders.h"
#include "Platform.h"
#include <switchboard_onnx/export.h>

#include <map>
#include <vector>
#include <memory>

namespace switchboard::extensions::onnx {

/**
 *  ModelInstance class.
 *  @brief Interact with Onnx models
 */
class SWITCHBOARDONNX_EXPORT ModelInstance {
public:
    ModelInstance(const OnnxProvider& provider = CpuProvider{});
    ~ModelInstance();

    /**
     *  @brief Run inference on provided input and predict the result in output
     *  this method is currently not used but can be useful
     *
     *  @param input Input vector
     *  @param output Output vector
     */
    template <typename T, typename U>
    void process(std::vector<T>& input, std::vector<U>& output);

    /**
     *  @brief Set data for a model input
     *
     *  @param name Input name
     *  @param data Vector<T> containing input data
     */
    template <typename T>
    void setInputByRef(const std::string& name, std::vector<T>& data);

    /**
     *  @brief Get data from a model onput
     *
     *  @param name Output name
     *  @return Vector<T> containing output data
     */
    template <typename T>
    std::vector<T> getOutputValues(const std::string& name);

    /**
     *  @brief Reset data for input
     *
     *  @param name Input name
     */
    template <typename T>
    void resetInput(const std::string& name);

    /**
     *  @brief Resize input
     *
     *  @param name Input name
     *  @param size Input size
     */
    template <typename T>
    void resizeInput(const std::string& name, int size);

    template <typename T>
    void resetOutput(const std::string& name);

    /**
     *  @brief Set dimension of input
     *
     *  @param name Input name
     *  @param index Index of the input's shape array where dimension size will be set
     *  @param dim Input vector containing shape information
     */
    void setDimAtIndex(const std::string& name, int index, int64_t dim);

    /**
     *  @brief Set shape of input tensor
     *
     *  @param name Input name
     *  @param dims Input vector containing shape information
     */
    void setDims(const std::string& name, std::vector<int64_t> dims);

    /**
     *  @brief Assigns data from a model output to a model input
     *  @param outputName Name of the output to assign data from
     *  @param inputName Name of the input to assign data output to
     */
    void bindIO(const std::string& outputName, const std::string& inputName) const;

    /**
     *  @brief Runs inference on model inputs
     */
    void runPrediction() const;


    /**
     *  @brief Get list of inputs of the loaded model
     *
     *  @return List of TensorInfo containing all inputs of the model
     */
    std::vector<TensorInfo> getModelInputInfo() const;

    /**
     *  @brief Get list of outputs of the loaded model
     *
     *  @return List of TensorInfo containing all inputs of the model
     */
    std::vector<TensorInfo> getModelOutputInfo() const;

    /**
     *  @brief Loads Model via string path
     *
     *  @param path String pointer to the model
     */
    void loadModelFromPath(const std::string& path) const;

    /**
     *  @brief Loads Model via pointer to model data
     *
     *  @param data Pointer to the model data
     *  @param size Size of model data in bytes
     */
    void loadModelFromData(const char* data, const unsigned int size) const;

    /**
     *  @brief Check if model has been loaded
     *
     *  @return True if model is loaded, false otherwise
     */
    bool isModelLoaded() const;

private:
    void populateModelInfo() const;

    class Impl;
    std::unique_ptr<Impl> pImpl;
};

}
