// Copyright © 2024 Intel Corporation
// SPDX-License-Identifier: Apache 2.0
// LEGAL NOTICE: Your use of this software and any required dependent software (the “Software Package”)
// is subject to the terms and conditions of the software license agreements for the Software Package,
// which may also include notices, disclaimers, or license terms for third party or open source software
// included in or with the Software Package, and your use indicates your acceptance of all such terms.
// Please refer to the “third-party-programs.txt” or other similarly-named text file included with the
// Software Package for additional details.

#ifndef VPUNN_RUNTIME_H
#define VPUNN_RUNTIME_H

///@todo: proposal to rename file into runtime_nn

#include <string>
#include "core/profiling.h"
#include "inference/inference_execution_data.h"
#include "inference/model.h"
#include "inference/model_version.h"

/// @brief top namespace for VPUNN cost model library
namespace VPUNN {

/**
 * @brief VPUNN runtime model
 *
 */
class Runtime {
private:
    const InferenceModel model;  ///< the NN loaded from a file/buffer (flatbuffer)
    // InferenceExecutionData model_buffer_data;  ///< the memory/buffers used for executing a model (in/out and inter
    ///< layer buffers). It is paired with the model at creation.

    const bool profile;
    ModelVersion model_version;  ///< holds the version info about the loaded model

public:
    /**
     * @brief Construct a new Runtime object
     *
     * @param filename .vpunn model
     * @param batch model batch size
     * @param profile enable/disable profiling
     */
    explicit Runtime(const std::string& filename, bool profile = false)
            : model(filename.c_str()),
              // model_buffer_data(createNewInferenceExecutionData(batch)),
              profile(profile),
              model_version() {
        if (initialized()) {
            model_version.parse_name(model.network_name());
        }
    }

    /**  @brief provides version info for loaded model
     * The provided reference is always up to date with the loaded model
     * The info are conditioned(otherwise default) by the successful loading of the model.
     *
     * @return a long lived reference to the version information
     */
    const ModelVersion& model_version_info() const {
        return model_version;
    }

    /**
     * @brief Construct a new Runtime object
     *
     * @param model_data .vpunn model buffer
     * @param model_data_length buffer size
     * @param copy_model_data enable/disable memcopy of the module buffer
     * @param batch model batch size
     * @param profile enable/disable profiling
     */
    explicit Runtime(const char* model_data, size_t model_data_length, bool copy_model_data, bool profile = false)
            : model(model_data, model_data_length, copy_model_data),
              // model_buffer_data(createNewInferenceExecutionData(batch)),
              profile(profile),
              model_version() {
        if (initialized()) {
            model_version.parse_name(model.network_name());
        }
    }

    /**
     * @brief Check if the NN model is initialized
     *
     * @return true if the NN model is initialized
     * @return false if the NN model is not initialized
     */
    bool initialized() const {
        return model.is_initialized();
    }

    InferenceExecutionData createNewInferenceExecutionData(const unsigned int batch) const {
        return InferenceExecutionData(batch, model.get_model());  // RVO
    }

private:
    ///**
    // * @brief Get the model input tensors
    // *
    // * @return std::vector<Tensor<float>*>
    // */
    // std::vector<const Tensor<float>*> input_tensors() const {
    //    return model_buffer_data.input_tensors();
    //}

    ///**
    // * @brief Get the model output tensors
    // *
    // * @return std::vector<Tensor<float>*>
    // */
    // std::vector<const Tensor<float>*> output_tensors() const {
    //    return model_buffer_data.output_tensors();
    //}

public:
    ///**
    // * @brief Get the model input tensors shapes
    // *
    // * @return std::vector<std::vector<unsigned int>>
    // */
    // std::vector<std::vector<unsigned int>> input_shapes() const {
    //    std::vector<std::vector<unsigned int>> shapes;
    //    for (const auto& tensor : input_tensors()) {
    //        shapes.push_back(tensor->shape());
    //    }
    //    return shapes;
    //}

    ///**
    // * @brief Get the model output tensors shapes
    // *
    // * @return std::vector<std::vector<unsigned int>>
    // */
    // std::vector<std::vector<unsigned int>> output_shapes() const {
    //    std::vector<std::vector<unsigned int>> shapes;
    //    for (const auto& tensor : output_tensors()) {
    //        shapes.push_back(tensor->shape());
    //    }
    //    return shapes;
    //}

    /**
     * @brief Run the inference
     *
     * @tparam T input input_array datatype
     * @param input_array input data
     * @param input_size input data size
     * @return pointer to one or more values containing the inference result
     */
    // template <class T>
    // const T* predict(const T* input_array, const unsigned int input_size) {
    //     model_buffer_data.set_inputs(input_array, input_size);  // might throw if input mismatch
    //     const auto t1{profile ? tick() : no_tick()};
    //     model.predict(model_buffer_data);
    //     if (profile) {
    //         const auto delta = tock(t1);
    //         Logger::info() << "Execution time: " << delta << " ms";
    //     }
    //     return model_buffer_data.get_outputs<T>();
    // }

    /**
     * @brief Run the inference in an external buffer
     *
     * @tparam T input input_array datatype
     * @param input_array input data
     * @param input_size input data size
     * @return pointer to one or more values containing the inference result, it is a pointer inside of thr
     * InferenceExecutionData buffer
     */
    template <class T>
    const T* predict(const T* input_array, const unsigned int input_size,
                     InferenceExecutionData& current_model_buffer_data) const {
        current_model_buffer_data.set_inputs(input_array,
                                             input_size);  // might throw if input mismatch
        const auto t1{profile ? tick() : no_tick()};
        model.predict(current_model_buffer_data);
        if (profile) {
            const auto delta = tock(t1);
            Logger::info() << "Execution time: " << delta << " ms";
        }
        return current_model_buffer_data.get_outputs<T>();
    }

    /**
     * @brief Run the inference
     *
     * @tparam T input input_array datatype
     * @param input_tensor input data
     * @return one or more values containing the inference result
     */
    // template <class T>
    // const std::vector<T> predict(const std::vector<T>& input_tensor) {
    //     model_buffer_data.set_inputs(input_tensor.data(),
    //                                  static_cast<unsigned int>(input_tensor.size()));  // might throw if input
    //                                  mismatch
    //     const auto t1{profile ? tick() : no_tick()};
    //     model.predict(model_buffer_data);
    //     if (profile) {
    //         const auto delta = tock(t1);
    //         Logger::info() << "Execution time: " << delta << " ms";
    //     }
    //     return model_buffer_data.get_outputs_copy_as_vector<T>();
    // }

    template <class T>
    const std::vector<T> predict(const std::vector<T>& input_tensor,
                                 InferenceExecutionData& current_model_buffer_data) const {
        current_model_buffer_data.set_inputs(
                input_tensor.data(),
                static_cast<unsigned int>(input_tensor.size()));  // might throw if input mismatch
        const auto t1{profile ? tick() : no_tick()};
        model.predict(current_model_buffer_data);
        if (profile) {
            const auto delta = tock(t1);
            Logger::info() << "Execution time: " << delta << " ms";
        }
        return current_model_buffer_data.get_outputs_copy_as_vector<T>();
    }

    /// just a wrapper for the model prediction. The buffer has to be with input data  and will contain output data
    void predict(InferenceExecutionData& current_model_buffer_data) const {
        const auto t1{profile ? tick() : no_tick()};
        model.predict(current_model_buffer_data);
        if (profile) {
            const auto delta = tock(t1);
            Logger::info() << "Execution time: " << delta << " ms";
        }
    }
};

}  // namespace VPUNN

#endif  // VPUNN_H
