#include <time.h>
#include <thread>
#include <iostream>
#include <stdlib.h>
#include <dirent.h>
#include <unistd.h>
#include <pthread.h>
#include <bits/stdint-uintn.h>
#include "acllite/AclLiteUtils.h"
#include "acllite/AclLiteResource.h"
#include "queue.h"
#include "acl/acl.h"
#include "object_detect.h"
#include "acllite/AclLiteUtils.h"
#include "acllite/AclLiteResource.h"
#include <vector>
#include <fstream>
#include <chrono>


const uint32_t deviceIdList [4]= {0,1,2,3};
uint32_t modelIdList[4];
const int32_t deviceNum = 4;
aclrtContext context;
std::vector<aclrtContext> cxtGroup;
void* pictureHostData = nullptr;
std::vector<void*> pictureDevDataGroup = {nullptr, nullptr, nullptr, nullptr};
std::vector<void*> outputDevDataGroup = {nullptr, nullptr, nullptr, nullptr};
std::vector<float> timeList;



int32_t ReadPictureToHost(const char *picturePath)
{
    int32_t pictureDataSize = 0;
    std::string fileName = picturePath;
    std::ifstream binFile(fileName, std::ios::binary);
    binFile.seekg(0, binFile.end);
    pictureDataSize = binFile.tellg();
    binFile.seekg(0, binFile.beg);
    aclError ret = aclrtMallocHost(&pictureHostData, pictureDataSize);
    if(ret!=ACL_SUCCESS){
        printf("malloc host failed! \n");
    }
    binFile.read((char*)pictureHostData, pictureDataSize);
    binFile.close();
    return pictureDataSize;
}


void CopyDataFromHostToDevice(void* pictureHostData, void* pictureDeviceData, const uint32_t pictureDataSize)
{
    aclError ret = aclrtMalloc(&pictureDeviceData, pictureDataSize, ACL_MEM_MALLOC_HUGE_FIRST);
    if(ret != ACL_SUCCESS){
        printf("malloc on device failed!");
    }
    ret = aclrtMemcpy(pictureDeviceData, pictureDataSize, pictureHostData, pictureDataSize, ACL_MEMCPY_HOST_TO_DEVICE);
    if(ret != ACL_SUCCESS){
        printf("memory copy failed!");
    }
    std::cout << "------------------------------------------" << std::endl;
    std::cout << "host data: " << pictureHostData << std::endl;
    std::cout << "device data: " << pictureDeviceData << std::endl;
}


aclmdlDataset* CreateModelInput(void* pictureDeviceData, const uint32_t pictureDataSize)
{
    // 创建aclmdlDataset类型的数据，描述模型推理的输入
    auto inputDataSet = aclmdlCreateDataset();
    auto inputDataBuffer = aclCreateDataBuffer(pictureDeviceData, pictureDataSize);  // TODO: how to delete
    aclError ret = aclmdlAddDatasetBuffer(inputDataSet, inputDataBuffer);
    if(ret != ACL_SUCCESS){
        printf("create input dataset failed!");
    }
    return inputDataSet;
}


aclmdlDataset* CreateModelOutput(const uint32_t modelId, void* outputDeviceData)
{
    auto modelDesc =  aclmdlCreateDesc();
    aclmdlGetDesc(modelDesc, modelId);

    auto outputDataSet = aclmdlCreateDataset();
    auto outputDataSize = aclmdlGetOutputSizeByIndex(modelDesc, 0);
    auto inputDataSize = aclmdlGetInputSizeByIndex(modelDesc, 0);
    // std::cout << "input data size: " << inputDataSize << std::endl;
    // std::cout << "output data size: " << outputDataSize << std::endl;

    aclrtMalloc(&outputDeviceData, outputDataSize, ACL_MEM_MALLOC_HUGE_FIRST);
    auto outputDataBuffer = aclCreateDataBuffer(outputDeviceData, outputDataSize); // TODO: how to delete
    aclmdlAddDatasetBuffer(outputDataSet, outputDataBuffer);

    // aclmdlDestroyDesc(modelDesc);
    // aclDestroyDataBuffer(outputDataBuffer);
    return outputDataSet;
}



void ModelExecute(const aclrtContext context,  const uint32_t modelId, const aclmdlDataset* input, aclmdlDataset* output){
    for(uint32_t i = 0; i<1; ++i){
        aclError ret = aclrtSetCurrentContext(context);
        if( ret != ACL_SUCCESS){
            printf("set context failed!");
        }
        auto time1  = std::chrono::steady_clock::now();
        std::cout << "now running model : " << modelId << std::endl;
        ret = aclmdlExecute(modelId, input, output);    
        if(ret != ACL_SUCCESS){
            printf("execute model %d failed! error code is %d \n", modelId, ret);
        }
        auto time2 = std::chrono::steady_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(time2 - time1);
        //std::cout << "duration inside execute: " << duration.count() << std::endl;
    }
}



int main(){
    aclInit(nullptr);
    aclError ret;

    for(int32_t i =0;i<deviceNum; ++i){
        ret = aclrtSetDevice(i);
        if(ret != ACL_SUCCESS){
            printf("set device %d failed! \n", i);
         }
        ret = aclrtCreateContext(&context, i);
        if(ret != ACL_SUCCESS){
           printf("create context %d failed! \n", i);
        }
        cxtGroup.push_back(context);
    }

    const char* picturePath = "../data/dog1_1024_683.bin";
    auto pictureDataSize  = ReadPictureToHost(picturePath);

    for(int32_t i=0; i<deviceNum; ++i){
        ret = aclrtSetCurrentContext(cxtGroup[i]);
        if(ret != ACL_SUCCESS){
            printf("set context %d failed!", i);
        }
        CopyDataFromHostToDevice(pictureHostData, pictureDevDataGroup[i], pictureDataSize);
    }

    const char* modelPath = "../model/resnet50.om"; 
    // size_t work_size;
    // size_t weight_size;
    // aclmdlQuerySize(modelPath, &work_size, &weight_size);
    // std::cout << "model work size: " << work_size << std::endl;
    // std::cout <<"model weight size:" << weight_size << std::endl;
    
    std::vector<aclmdlDataset*> inputDataSetGroup;
    std::vector<aclmdlDataset*> outputDataSetGroup;
    for(int32_t i =0;i<deviceNum; ++i){
        ret = aclrtSetCurrentContext(cxtGroup[i]);
        if(ret != ACL_SUCCESS){
            printf("set context %d failed!", i);
        }
        ret = aclmdlLoadFromFile(modelPath, &(modelIdList[i]));
        if(ret != ACL_SUCCESS){
            printf("load model %d failed!", i);
        }
        
        auto input = CreateModelInput(pictureDevDataGroup[i], pictureDataSize);
        auto output = CreateModelOutput(modelIdList[i], outputDevDataGroup[i]);
        inputDataSetGroup.push_back(input);
        outputDataSetGroup.push_back(output);
    }

    
    // std::thread task0(ModelExecute, std::ref(cxtGroup[0]), std::ref(modelIdList[0]), std::ref(inputDataSetGroup[0]), std::ref(outputDataSetGroup[0]));
    // std::thread task1(ModelExecute, std::ref(cxtGroup[1]), std::ref(modelIdList[1]), std::ref(inputDataSetGroup[1]), std::ref(outputDataSetGroup[1]));
    // std::thread task2(ModelExecute, std::ref(cxtGroup[2]), std::ref(modelIdList[2]), std::ref(inputDataSetGroup[2]), std::ref(outputDataSetGroup[2]));
    // std::thread task3(ModelExecute, std::ref(cxtGroup[3]), std::ref(modelIdList[3]), std::ref(inputDataSetGroup[3]), std::ref(outputDataSetGroup[3]));

    
    // task0.join();
    // task1.join();
    // task2.join();
    // task3.join();

    ModelExecute(cxtGroup[0], modelIdList[0], inputDataSetGroup[0], outputDataSetGroup[0]); 

    

    aclrtFreeHost(pictureHostData);
    for(int32_t i=0;i<deviceNum;++i){
        aclrtFree(pictureDevDataGroup[i]);
        aclmdlDestroyDataset(inputDataSetGroup[i]);
        aclmdlDestroyDataset(outputDataSetGroup[i]);
        aclmdlUnload(modelIdList[i]);
        aclrtDestroyContext(cxtGroup[i]);
    }
    aclFinalize();


}