/***********************************************************************************
MIT License

Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice (including the next paragraph) shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 ************************************************************************************/
#include <assert.h>
#include <onnxruntime_cxx_api.h>

#include <iostream>
#include <sstream>
#include <vector>
#include <codecvt>
#include <filesystem>

#include "npu_util.h"


static int get_num_elements(const std::vector<int64_t>& v) {
    int total = 1;
    for (auto& i : v)
        total *= (int)i;
    return total;
}

template <typename T>
std::ostream& operator<<(std::ostream& os, const std::vector<T>& v)
{
    os << "[";
    for (int i = 0; i < v.size(); ++i)
    {
        os << v[i];
        if (i != v.size() - 1)
        {
            os << ", ";
        }
    }
    os << "]";
    return os;
}

// pretty prints a shape dimension vector
static std::string print_shape(const std::vector<int64_t>& v) {
    std::stringstream ss("");
    for (size_t i = 0; i < v.size() - 1; i++)
        ss << v[i] << "x";
    ss << v[v.size() - 1];
    return ss.str();
}

static std::string print_tensor(Ort::Value& tensor) {
    auto shape = tensor.GetTensorTypeAndShapeInfo().GetShape();
    auto nelem = get_num_elements(shape);
    auto tensor_ptr = tensor.GetTensorMutableData<float>();

    std::stringstream ss("");
    for (auto i = 0; i < nelem; i++)
        ss << tensor_ptr[i] << " ";
    return ss.str();
}

template <typename T>
Ort::Value vec_to_tensor(std::vector<T>& data, const std::vector<std::int64_t>& shape) {
  Ort::MemoryInfo mem_info =
      Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
  auto tensor = Ort::Value::CreateTensor<T>(mem_info, data.data(), data.size(), shape.data(), shape.size());
  return tensor;
}

std::string get_program_dir()
{
    char* exe_path; _get_pgmptr(&exe_path); // full path and name of the executable
    return std::filesystem::path(exe_path).parent_path().string(); // directory in which the executable is located
}


int runtest(std::string& model_name, std::unordered_map<std::string, std::string>& vai_ep_options)
{    
    int64_t batch_size  = 1;

    printf("Creating ORT env\n");
    Ort::Env env(ORT_LOGGING_LEVEL_ERROR, "quicktest");

    printf("Initializing session options\n");
    auto session_options = Ort::SessionOptions();

    if (vai_ep_options.empty()==false) // If VAI EP options are provided, initialize the VitisAI EP
    {
        printf("Configuring VAI EP\n");
        try {
            session_options.AppendExecutionProvider_VitisAI(vai_ep_options);
        }
        catch (const std::exception& e) {
            std::cerr << "Exception occurred in appending execution provider: " << e.what() << std::endl;
        }
    }

    printf("Creating ONNX Session\n");
    auto session = Ort::Session(env, std::basic_string<ORTCHAR_T>(model_name.begin(), model_name.end()).c_str(), session_options);

    // Get names and shapes of model inputs and outputs
    Ort::AllocatorWithDefaultOptions allocator;
    auto input_count       = session.GetInputCount();
    auto input_names       = std::vector<std::string>();
    auto input_names_char  = std::vector<const char*>();
    auto input_shapes      = std::vector<std::vector<int64_t>>();
    auto output_count      = session.GetOutputCount();
    auto output_names      = std::vector<std::string>();
    auto output_names_char = std::vector<const char*>();
    auto output_shapes     = std::vector<std::vector<int64_t>>();
    for (size_t i = 0; i < input_count; i++)
    {
        auto shape = session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        std::string name = session.GetInputNameAllocated(i, allocator).get();
        input_names.emplace_back(name);
        input_names_char.emplace_back(input_names.at(i).c_str());
        input_shapes.emplace_back(shape);
    }
    for (size_t i = 0; i < output_count; i++)
    {
        auto shape = session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        std::string name = session.GetOutputNameAllocated(i, allocator).get();
        output_names.emplace_back(name);
        output_names_char.emplace_back(output_names.at(i).c_str());
        output_shapes.emplace_back(shape);
    }

    // Display model info
    std::cout << "ONNX model : " << model_name << std::endl;
    for (size_t i = 0; i < input_count; i++)
        std::cout << "  " << input_names.at(i) << " " << print_shape(input_shapes.at(i)) << std::endl;
    for (size_t i = 0; i < output_count; i++)
        std::cout << "  " << output_names.at(i) << " " << print_shape(output_shapes.at(i)) << std::endl;

    // The code which follows expects the model to have 1 input node and 1 output node.
    if (output_count != 1 && input_count != 1) {
        std::cout << "This version of the program only supports models with 1 input node and 1 output node. Exiting." << std::endl;
        exit(-1);
    }

    // If input shape has dynamic batch size, set it to a fixed value
    auto input_shape = input_shapes[0];
    if (input_shape[0] < 0) {
        std::cout << "Dynamic batch size detected. Setting batch size to " << batch_size << "." << std::endl;        
        input_shape[0] = batch_size;
    }

    printf("Running the model\n");
    for (int i = 0; i < 1; i++)
    {
        // Initialize input data with random numbers in the range [0, 255]
        std::vector<float> input_tensor_values(get_num_elements(input_shape));
        std::generate(input_tensor_values.begin(), input_tensor_values.end(), [&] { return (float)(rand() % 255); });

        // Initialize input tensor with input data
        std::vector<Ort::Value> input_tensors;
        input_tensors.emplace_back(vec_to_tensor<float>(input_tensor_values, input_shape));

        // Pass input tensors through model
        try {
            auto output_tensors = session.Run(
                    Ort::RunOptions(), 
                    input_names_char.data(), input_tensors.data(), input_names_char.size(), 
                    output_names_char.data(), output_names_char.size()
            );
            // std::cout << i << " : " << print_tensor(output_tensors[0]) << std::endl;
        }
        catch (const Ort::Exception& exception) {
            std::cout << "ERROR running model inference: " << exception.what() << std::endl;
            exit(-1);
        }
    }
    printf("-------------------------------------------------------\n");
    printf("Test PASSED!\n");
    printf("-------------------------------------------------------\n");
    printf("\n");
    
    return 0;
}


int run_on_cpu(std::string& model_name, std::string& exe_dir)
{
    // Leave VitisAI EP options empty to run on CPU
    std::unordered_map<std::string, std::string> vai_ep_options;

    // Full path to the ONNX model
    std::string model_path = exe_dir + "\\" + model_name;

    // Run test
    printf("-------------------------------------------------------\n");
    printf("Running quicktest on CPU                               \n");    
    printf("-------------------------------------------------------\n");
    return runtest(model_path, vai_ep_options);
}

int run_on_npu(std::string& model_name, std::string& exe_dir)
{
    printf("-------------------------------------------------------\n");
    printf("Performing compatibility check for VitisAI EP 1.4      \n");
    printf("-------------------------------------------------------\n");
    auto npu_info = npu_util::checkCompatibility_RAI_1_4();

    std::cout << " - NPU Device ID     : 0x" << std::hex << npu_info.device_id << std::dec << std::endl;
    std::cout << " - NPU Device Name   : " << npu_info.device_name << std::endl;
    std::cout << " - NPU Driver Version: " << npu_info.driver_version_string << std::endl;  
    switch (npu_info.check) {
        case npu_util::Status::OK:          
            std::cout << "Environment compatible for VitisAI EP" << std::endl;
            break;
        case npu_util::Status::NPU_UNRECOGNIZED:
            std::cout << "NPU type not recognized." << std::endl;
            std::cout << "Skipping run with VitisAI EP." << std::endl;
            return -1;           
            break;
        case npu_util::Status::DRIVER_TOO_OLD: 
            std::cout << "Installed drivers are too old." << std::endl;
            std::cout << "Skipping run with VitisAI EP." << std::endl;
            return -1;           
            break;
        case npu_util::Status::EP_TOO_OLD:
            std::cout << "VitisAI EP is too old." << std::endl;
            std::cout << "Skipping run with VitisAI EP." << std::endl;
            return -1;           
            break;
        default:
            std::cout << "Unknown state." << std::endl;
            std::cout << "Skipping run with VitisAI EP." << std::endl;
            return -1;           
            break;
    }
    std::cout << std::endl;

    // Set VitisAI EP options
    std::unordered_map<std::string, std::string> vai_ep_options;
    switch(npu_info.device_id) {
        case 0x1502: // PHX/HPT NPU
            vai_ep_options["cacheDir"] = exe_dir + "\\modelcache"; 
            vai_ep_options["cacheKey"] = "testmodel_phx"; 
            vai_ep_options["xclbin"]   = exe_dir + "\\xclbins\\phoenix\\1x4.xclbin";; 
            break;
        case 0x17F0: // STX/KRK NPU
            vai_ep_options["cacheDir"] = exe_dir + "\\modelcache"; 
            vai_ep_options["cacheKey"] = "testmodel_stx"; 
            vai_ep_options["xclbin"]   = exe_dir + "\\xclbins\\strix\\AMD_AIE2P_Nx4_Overlay.xclbin"; 
            break;
        default:
            std::cout << "Unsupported NPU device ID." << std::endl;
            return -1;
            break;
    }

    // Set environment variables
    _putenv("XLNX_VART_FIRMWARE=");           // Unset XLNX_VART_FIRMWARE (use VAI-EP option to set XCLBIN)
    _putenv("XLNX_TARGET_NAME=");             // Unset XLNX_TARGET_NAME   (rely on default value: AMD_AIE2P_Nx4_Overlay)

    // Full path to the ONNX model
    std::string model_path = exe_dir + "\\" + model_name;

    // Run test
    printf("-------------------------------------------------------\n");
    printf("Running quicktest on NPU                               \n");    
    printf("-------------------------------------------------------\n");
    return runtest(model_path, vai_ep_options);
}

int main(int argc, char* argv[]) 
{
    std::string exe_dir = get_program_dir();
    std::string model_name ="test_model.onnx";
 
    run_on_cpu(model_name, exe_dir);
    run_on_npu(model_name, exe_dir);

    return 0;
}