#include "onnxruntime_engine.h"
#include <iostream>
#include <chrono>

OnnxruntimeEngine::OnnxruntimeEngine(const std::string& modelName, bool useOpenvino) {
    const std::string logStr = "Trueface_SDK_v_onnxruntime"; // TODO Cyrus: Get the SDK version here using the static SDK method.
    m_envPtr = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, logStr.c_str()); // TODO Cyrus: The log level needs to be set based on config setting.

    m_sessionOpts.SetInterOpNumThreads(8); // TODO Cyrus: Configure this using the threads argument.
    m_sessionOpts.SetIntraOpNumThreads(8);

    // Sets graph optimization level
    // Available levels are
    // ORT_DISABLE_ALL -> To disable all optimizations
    // ORT_ENABLE_BASIC -> To enable basic optimizations (Such as redundant node removals)
    // ORT_ENABLE_EXTENDED -> To enable extended optimizations (Includes level 1 + more complex optimizations like node fusions)
    // ORT_ENABLE_ALL -> To Enable All possible opitmizations

    const auto providers = Ort::GetAvailableProviders();
    for (const auto& provider : providers) {
        std::cout << "Available providers: " << provider << std::endl;
    }

    if (useOpenvino) {
        OrtOpenVINOProviderOptions options;
        options.device_type = "CPU_FP32";
        options.num_of_threads = 8;
        m_sessionOpts.AppendExecutionProvider_OpenVINO(options);
        // For OpenVINO we disable all onnxruntime optimizations and allow OpenVINO to do it's own optimizations instead.
        m_sessionOpts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); // https://onnxruntime.ai/docs/execution-providers/OpenVINO-ExecutionProvider.html#cc-api-1
        std::cout << "Using OpenVINO execution provider" << std::endl;
    }
    else {
        m_sessionOpts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        //        m_sessionOpts.SetExecutionMode(ExecutionMode::ORT_PARALLEL);
        std::cout << "Using CPU execution provider" << std::endl;
    }

    // TODO Cyrus: Need to check that the model exist
    // TODO Cyrus: Need to decrypt the model here
    // TODO Cyrus: Think about if you want to support models which have not been encrypted.
    m_sessionPtr = std::make_unique<Ort::Session>(*m_envPtr, modelName.c_str(), m_sessionOpts);

    Ort::AllocatorWithDefaultOptions allocator;

    size_t numInputNodes = m_sessionPtr->GetInputCount();
    m_inputNames.reserve(numInputNodes);

    // TODO Cyrus: Should the following logs be based on TRACE level? Or not have them at all?
//    std::cout << "Number of inputs: " << numInputNodes << std::endl;

    for (int i = 0; i < numInputNodes; i++) {
        auto inputName = m_sessionPtr->GetInputNameAllocated(i, allocator);
        //        std::cout << "Input " << i << ", name = " << inputName.get() << std::endl;
        m_inputNames.push_back(_strdup(inputName.get()));

        auto typeInfo = m_sessionPtr->GetInputTypeInfo(i);
        auto tensorInfo = typeInfo.GetTensorTypeAndShapeInfo();

        auto type = tensorInfo.GetElementType();
        //        std::cout << "Input " << i << ", type = " << type << std::endl;

        m_inputDims = tensorInfo.GetShape();
        //        std::cout << "Input " << i << ", num_dims = " << m_inputDims.size() << std::endl;

        m_inputTensorSize = 1;
        for (int j = 0; j < m_inputDims.size(); j++) {
            if (m_inputDims[j] == -1) {
                m_inputDims[j] = 1;
            }
            //            std::cout << "Input " << i << ", dim " << j << " = " << m_inputDims[j] << std::endl;
            m_inputTensorSize *= m_inputDims[j];
        }
    }

    size_t numOutputNodes = m_sessionPtr->GetOutputCount();
    m_outputNames.reserve(numOutputNodes);

    //    std::cout << "Number of output = " << numOutputNodes << std::endl;

    for (int i = 0; i < numOutputNodes; i++) {
        auto outputName = m_sessionPtr->GetOutputNameAllocated(i, allocator);
        //        std::cout << "Output " << i << ", name = " << outputName.get() << std::endl;
        m_outputNames.push_back(_strdup(outputName.get()));

        auto typeInfo = m_sessionPtr->GetOutputTypeInfo(i);
        auto tensorInfo = typeInfo.GetTensorTypeAndShapeInfo();

        auto type = tensorInfo.GetElementType();
        //        std::cout << "Output " << i << ", type = " << type << std::endl;

        m_outputDims = tensorInfo.GetShape();
        //        std::cout << "Output " << i << ", num_dims = " << m_outputDims.size() << std::endl;

        m_outputTensorSize = 1;
        for (int j = 0; j < m_outputDims.size(); j++)
        {
            if (m_outputDims[j] == -1) {
                m_outputDims[j] = 1;
            }
            m_outputTensorSize *= m_outputDims[j];
            //            std::cout << "Output " << i << ", dim " << j << " = " << m_outputDims[j] << std::endl;
        }
    }
}

OnnxruntimeEngine::~OnnxruntimeEngine() {
    m_sessionPtr->release();
}

std::vector<Ort::Value> OnnxruntimeEngine::getFeatureVector(const Ort::Value& inputTensor) {
    auto outputTensors = m_sessionPtr->Run(Ort::RunOptions{ nullptr }, m_inputNames.data(), &inputTensor, 1, m_outputNames.data(), 1);
    return std::move(outputTensors);
}

unsigned long OnnxruntimeEngine::getInputTensorSize() const {
    return m_inputTensorSize;
}

std::vector<int64_t> OnnxruntimeEngine::getInputDimension() const {
    return m_inputDims;
}

std::vector<int64_t> OnnxruntimeEngine::getOutputDimension() const {
    return m_outputDims;
}

unsigned long OnnxruntimeEngine::getOutputTensorSize() const {
    return m_outputTensorSize;
}