/*
 * Copyright 2020 Xilinx Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef _N2CUBE_H_
#define _N2CUBE_H_


#include <errno.h>
#include <math.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <atomic>
#include <chrono>
#include <mutex>
#include <sstream>
#include <thread>
#include <chrono>
#include <iomanip>
#include <iostream>
#include <map>
#include <fstream>
#include <string>
#include <streambuf>
#include <unordered_map>
#include <future>
#include <map>
#include <list>
#include <iterator>

#include <pthread.h>
#include <semaphore.h>
#include <future>
#include <map>

#ifdef __cplusplus
extern "C" {
#endif

#include <stdint.h>

#define VITIS_AI_VERSION    "1.3"
#define N2Cube_VERSION   "4.3"

/* DPU Task runtime mode definitions */

/* Task in normal mode (defaul mode) */
#define T_MODE_NORMAL        (0)

/* Task in profiling mode in order to collect performance stastics for each DPU Node */
#define T_MODE_PROFILE       (1<<0)

/* Task in debug mode in order to dump each Node's Code/Bias/Weights/Input/Output raw data for debugging */
#define T_MODE_DEBUG         (1<<1)

/* Exported data structures of DPU Kernel/Task/Tensor */
struct  dpu_kernel;
struct  dpu_task;
struct  task_tensor;

typedef struct dpu_kernel  DPUKernel;
typedef struct dpu_task    DPUTask;
typedef struct task_tensor DPUTensor;

/* The exception handling mode */
#define N2CUBE_EXCEPTION_MODE_PRINT_AND_EXIT 0
#define N2CUBE_EXCEPTION_MODE_RET_ERR_CODE 1
int dpuSetExceptionMode(int mode);
int dpuGetExceptionMode();
const char *dpuGetExceptionMessage(int error_code);

/* Open & initialize the usage of DPU device */
int dpuOpen();

/* Close & finalize the usage of DPU device */
int dpuClose();

#ifdef __cplusplus
}
#endif

/* input image with neon */
void dpuProcessNormalizion(int8_t* dst, uint8_t* src, int h, int w, float* shift, float scale, int stride);

/* Get mean from model 
   This API is only available for Caffe model*/
int dpuGetKernelMean(DPUTask *task, float *mean, int count);

/* Load a DPU Kernel and allocate DPU memory space for
   its Code/Weight/Bias segments */
DPUKernel *dpuLoadKernel(const char *netName);

/* Get the total size in memory of all the inputs of the network,
   which involve the interleaves between two inputs. */
int dpuGetInputTotalSize(DPUTask *task);

/* Get the total size in memory of all the outputs of the network,
   which involve the interleaves between two outputs. */
int dpuGetOutputTotalSize(DPUTask *task);

/* Allocate DPU memory space with 'size' length.
   It will return a device handler. See dpuFreeMem also.
   Just used for split_io sample. */
void* dpuAllocMem(int size, int8_t* &addrVirtual, int8_t* &addrPhysical);

/* Release the memory allocated through API 'dpuAllocMem'.
   Pass the 'handler' returned from API 'dpuAllocMem'.
   Just used for split_io sample. */
void dpuFreeMem(void* handler);

/* Sync mem from host to DPU, be called after set input.
   Pass the 'handler' returned from API 'dpuAllocMem'.
   Just used for split_io sample. */
void dpuSyncMemToDev(void *handler, uint32_t offset, uint32_t size);

/* Sync mem from DPU to host, be called after output generated by DPU.
   Pass the 'handler' returned from API 'dpuAllocMem'.
   Just used for split_io sample. */
void dpuSyncDevToMem(void *handler, uint32_t offset, uint32_t size);

/* Set mean values for DPU Kernel */
int dpuSetKernelMeanValue(DPUKernel *kernel, int mean1, int mean2, int mean3);

/* Destroy a DPU Kernel and release its associated resources */
int dpuDestroyKernel(DPUKernel *kernel);

/* Instantiate a DPU Task from one DPU Kernel, allocate its private
   working memory buffer and prepare for its execution context */
DPUTask *dpuCreateTask(DPUKernel *kernel, int mode);

/* Bind base address for all the input tensors of a network.
   'addrVirtual' is the address from host side,
   'addrPhysical' is the address from DPU side. */
int dpuBindInputTensorBaseAddress(DPUTask *task, int8_t *addrVirtual, int8_t *addrPhysical);

/* Bind base address for all the output tensors of a network.
   'addrVirtual' is the address from host side,
   'addrPhysical' is the address from DPU side. */
int dpuBindOutputTensorBaseAddress(DPUTask *task, int8_t *addrVirtual, int8_t *addrPhysical);

/* Set or get the priority of one DPU task. Priority range is 0 to 15, 0 has the highest priority.
   The priority of the task when it was created defaults to 15. */
int dpuSetTaskPriority(DPUTask *task, uint8_t priority);
uint8_t dpuGetTaskPriority(DPUTask *task);

/* Set or get the core mask of the task binding. Each bit represents a core.
   The default value of coreMask is 0x7 when the number of cores is 3. */
int dpuSetTaskAffinity(DPUTask *task, uint32_t coreMask);
uint32_t dpuGetTaskAffinity(DPUTask *task);

/* Launch the running of DPU Task */
int dpuRunTask(DPUTask *task);

/* Remove a DPU Task, release its working memory buffer and destroy
   associated execution context */
int dpuDestroyTask(DPUTask *task);

/* Enable dump facility of DPU Task while running for debugging purpose */
int dpuEnableTaskDebug(DPUTask *task);

/* Enable profiling facility of DPU Task while running to get its performance metrics */
int dpuEnableTaskProfile(DPUTask *task);

/* Get the execution time of DPU Task */
long long dpuGetTaskProfile(DPUTask *task);

/* Get the execution time of DPU Node */
long long dpuGetNodeProfile(DPUTask *task, const char*nodeName);


/*
 * API for both single IO and multiple IO.
 * For multiply IO, should specify the input/output tensor idx.
 */

/* Get total number of input Tensor of DPU Task */
int dpuGetInputTensorCnt(DPUTask * task, const char * nodeName);

/* Get input Tensor of DPU Task */
DPUTensor* dpuGetInputTensor(DPUTask *task, const char*nodeName, int idx = 0);

/* Get the start address of DPU Task's input Tensor */
int8_t* dpuGetInputTensorAddress(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the size (in byte) of one DPU Task's input Tensor */
int dpuGetInputTensorSize(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the scale value (DPU INT8 quantization) of one DPU Task's input Tensor */
float dpuGetInputTensorScale(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the height dimension of one DPU Task's input Tensor */
int dpuGetInputTensorHeight(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the width dimension of one DPU Task's input Tensor */
int dpuGetInputTensorWidth(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the channel dimension of one DPU Task's input Tensor */
int dpuGetInputTensorChannel(DPUTask *task, const char *nodeName, int idx = 0);

/* Get total number of output Tensor of DPU Task */
int dpuGetOutputTensorCnt(DPUTask * task, const char * nodeName);

/* Get output Tensor of one DPU Task */
DPUTensor* dpuGetOutputTensor(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the start address of one DPU Task's output Tensor */
int8_t* dpuGetOutputTensorAddress(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the size (in byte) of one DPU Task's output Tensor */
int dpuGetOutputTensorSize(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the scale value (DPU INT8 quantization) of one DPU Task's output Tensor */
float dpuGetOutputTensorScale(DPUTask *task, const char *nodeName, int idx = 0);

/* Get the height dimension of one DPU Task's output Tensor */
int dpuGetOutputTensorHeight(DPUTask *task, const char *nodeName, int idx = 0);

/*  Get the channel dimension of one DPU Task's output Tensor */
int dpuGetOutputTensorWidth(DPUTask *task, const char *nodeName, int idx = 0);

/* Get DPU Node's output tensor's channel */
int dpuGetOutputTensorChannel(DPUTask *task, const char *nodeName, int idx = 0);

/* Set DPU Task's input Tensor with data stored under Caffe
   Blob's order (channel/height/width) in INT8 format */
int dpuSetInputTensorInCHWInt8(DPUTask *task, const char *nodeName, int8_t *data, int size, int idx = 0);

/* Set DPU Task's input Tensor with data stored under Caffe
   Blob's order (channel/height/width) in FP32 format */
int dpuSetInputTensorInCHWFP32(DPUTask *task, const char *nodeName, float *data, int size, int idx = 0);

/* Set DPU Task's input Tensor with data stored under DPU
   Tensor's order (height/width/channel) in INT8 format */
int dpuSetInputTensorInHWCInt8(DPUTask *task, const char *nodeName, int8_t *data, int size, int idx = 0);

/* Set DPU Task's input Tensor with data stored under DPU
   Tensor's order (height/width/channel) in FP32 format */
int dpuSetInputTensorInHWCFP32(DPUTask *task, const char *nodeName, float *data, int size, int idx = 0);

/* Get DPU Task's output Tensor and store them under Caffe
   Blob's order (channel/height/width) in INT8 format */
int dpuGetOutputTensorInCHWInt8(DPUTask *task, const char *nodeName, int8_t *data, int size, int idx = 0);

/* Get DPU Task's output Tensor and store them under Caffe
   Blob's order (channel/height/width) in FP32 format */
int dpuGetOutputTensorInCHWFP32(DPUTask *task, const char *nodeName, float *data, int size, int idx = 0);

/* Get DPU Task's output Tensor and store them under DPU
   Tensor's order (height/width/channel) in INT8 format */
int dpuGetOutputTensorInHWCInt8(DPUTask *task, const char *nodeName, int8_t *data, int size, int idx = 0);

/* Get DPU Task's output Tensor and store them under DPU
   Tensor's order (height/width/channel) in FP32 format */
int dpuGetOutputTensorInHWCFP32(DPUTask *task, const char *nodeName, float *buffer, int size, int idx = 0);

DPUTensor *dpuGetBoundaryIOTensor(DPUTask *task, const char *tensorName);

/* Get the size of one DPU Tensor */
int dpuGetTensorSize(DPUTensor* tensor);

/* Get the start address of one DPU Tensor */
int8_t* dpuGetTensorAddress(DPUTensor* tensor);

/* Get the scale value of one DPU Tensor */
float dpuGetTensorScale(DPUTensor* tensor);

/* Get the height dimension of one DPU Tensor */
int dpuGetTensorHeight(DPUTensor* tensor);

/* Get the width dimension of one DPU Tensor */
int dpuGetTensorWidth(DPUTensor* tensor);

/* Get the channel dimension of one DPU Tensor */
int dpuGetTensorChannel(DPUTensor* tensor);

/* Compute softmax */
int dpuRunSoftmax(int8_t *input, float *output, int numClasses, int batchSize, float scale);

#endif
