#include <signal.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <chrono>
#include <array>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <cstring>
#include <vector>
#include <dirent.h>

#include "sample_log.h"
#include "ax_interpreter_external_api.h"
#include "ax_sys_api.h"
#include "joint.h"
#include "joint_adv.h"
#include "ax_npu_imgproc.h"
#include "joint.h"
#include "joint_adv.h"

#include "detection.hpp"
#include "sample_run_joint_api.h"

std::vector<std::string> KEYPOINT_CLASS_NAMES = {
    "person", "hjjhand"};
std::vector<std::string> SEG_CLASS_NAMES = {
    "shoe", "xxx"};

namespace
{
    const float PROB_THRESHOLD = 0.5f;
    const float NMS_THRESHOLD = 0.45f;

    typedef struct _ax_cv_image
    {
        int h;
        int w;
        int stride_w;
        AX_NPU_CV_FrameDataType color_space;
    } ax_cv_image;

    typedef struct sampleRunJoint_PRIV_CTX_S
    {
        AX_JOINT_HANDLE pHandle;
        AX_JOINT_IO_SETTING_T stJointIoSetting;
        AX_JOINT_EXECUTION_CONTEXT stJointCtx;
        AX_JOINT_IO_T stJointIoArr;
        AX_JOINT_IO_BUFFER_T *pstIOBuffer;
        std::array<int, 2> arrInputSize; // [0]: height; [1]: width;
        AX_NPU_CV_FrameDataType eDtype;
        AX_NPU_CV_Image *pDstNV12;
        // 1=检测+关键点 2=分割
        AX_S32 modelType;
    } SampleRunJoint_PRIV_CTX_S;

    static SampleRunJoint_PRIV_CTX_S g_stInstCtx = {0};

    bool read_file(const char *fn, std::vector<char> &data)
    {
        FILE *fp = fopen(fn, "r");
        if (fp != nullptr)
        {
            fseek(fp, 0L, SEEK_END);
            auto len = ftell(fp);
            fseek(fp, 0, SEEK_SET);
            data.clear();
            size_t read_size = 0;
            if (len > 0)
            {
                data.resize(len);
                read_size = fread(data.data(), 1, len, fp);
            }
            fclose(fp);
            return read_size == (size_t)len;
        }
        return false;
    }

    void brief_io_info(const AX_JOINT_IO_INFO_T *io_info)
    {
        auto describe_shape_type = [](AX_JOINT_TENSOR_LAYOUT_T type) -> const char *
        {
            switch (type)
            {
            case AX_JOINT_TENSOR_LAYOUT_T::JOINT_TENSOR_LAYOUT_NHWC:
                return "NHWC";
            case AX_JOINT_TENSOR_LAYOUT_T::JOINT_TENSOR_LAYOUT_NCHW:
                return "NCHW";
            default:
                return "unknown";
            }
        };
        auto describe_data_type = [](AX_JOINT_DATA_TYPE_T type) -> const char *
        {
            switch (type)
            {
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT8:
                return "uint8";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT16:
                return "uint16";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_FLOAT32:
                return "float32";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT16:
                return "sint16";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT8:
                return "sint8";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT32:
                return "sint32";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT32:
                return "uint32";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_FLOAT64:
                return "float64";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT64:
                return "sint64";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT64:
                return "uint64";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT10_PACKED:
                return "uint10_packed";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT12_PACKED:
                return "uint12_packed";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT14_PACKED:
                return "uint14_packed";
            case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT16_PACKED:
                return "uint16_packed";
            default:
                return "unknown";
            }
        };
        auto describe_memory_type = [](AX_JOINT_MEMORY_TYPE_T type) -> const char *
        {
            switch (type)
            {
            case AX_JOINT_MEMORY_TYPE_T::AX_JOINT_MT_PHYSICAL:
                return "Physical";
            case AX_JOINT_MEMORY_TYPE_T::AX_JOINT_MT_VIRTUAL:
                return "Virtual";
            default:
                return "unknown";
            }
        };
        auto describe_color_space = [](AX_JOINT_COLOR_SPACE_T cs) -> const char *
        {
            switch (cs)
            {
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_FEATUREMAP:
                return "FeatureMap";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_BGR:
                return "BGR";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RGB:
                return "RGB";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_GRAY:
                return "GRAY";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_NV12:
                return "NV12";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_NV21:
                return "NV21";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_YUV444:
                return "YUV444";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW8:
                return "RAW8";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW10:
                return "RAW10";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW12:
                return "RAW12";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW14:
                return "RAW14";
            case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW16:
                return "RAW16";
            default:
                return "unknown";
            }
        };
        printf("Max Batch Size %d\n", io_info->nMaxBatchSize);
        printf("Support Dynamic Batch? %s\n",
               io_info->bDynamicBatchSize == AX_TRUE ? "Yes" : "No");
        printf("Is FilterMode? %s\n\n",
               io_info->bFilterMode == AX_TRUE ? "Yes" : "No");
        if (io_info->eQuantizationType ==
            AX_JOINT_QUANTIZATION_TYPE_T::AX_JOINT_QT_BIT_MIX)
        {
            printf("Quantization Type is mix\n\n");
        }
        else if (io_info->eQuantizationType ==
                 AX_JOINT_QUANTIZATION_TYPE_T::AX_JOINT_QT_BIT_8)
        {
            printf("Quantization Type is 8 bit\n\n");
        }
        else if (io_info->eQuantizationType ==
                 AX_JOINT_QUANTIZATION_TYPE_T::AX_JOINT_QT_BIT_16)
        {
            printf("Quantization Type is 16 bit\n\n");
        }
        else if (io_info->eQuantizationType ==
                 AX_JOINT_QUANTIZATION_TYPE_T::AX_JOINT_QT_BIT_32)
        {
            printf("Quantization Type is 32 bit\n\n");
        }
        else
        {
            printf("Quantization Type is unknown\n\n");
        }

        for (uint32_t i = 0; i < io_info->nInputSize; ++i)
        {
            auto &input = io_info->pInputs[i];
            printf("Input[%d]: %s\n", i, input.pName);
            printf("    Shape [");
            for (uint32_t j = 0; j < input.nShapeSize; ++j)
            {
                printf("%d", (int)input.pShape[j]);
                if (j + 1 < input.nShapeSize)
                    printf(", ");
            }
            printf("] %s %s %s %s\n", describe_shape_type(input.eLayout),
                   describe_data_type(input.eDataType),
                   input.pExtraMeta
                       ? describe_color_space(input.pExtraMeta->eColorSpace)
                       : "",
                   input.nQuantizationValue > 0
                       ? ("Q=" + std::to_string(input.nQuantizationValue)).c_str()
                       : "");
            printf("    Memory %s\n", describe_memory_type(input.eMemoryType));
            printf("    Size %u\n", input.nSize);
        }
        for (uint32_t i = 0; i < io_info->nOutputSize; ++i)
        {
            auto &output = io_info->pOutputs[i];
            printf("Output[%d]: %s\n", i, output.pName);
            printf("    Shape [");
            for (uint32_t j = 0; j < output.nShapeSize; ++j)
            {
                printf("%d", (int)output.pShape[j]);
                if (j + 1 < output.nShapeSize)
                    printf(", ");
            }
            printf("] %s %s %s\n", describe_shape_type(output.eLayout),
                   describe_data_type(output.eDataType),
                   output.nQuantizationValue > 0
                       ? ("Q=" + std::to_string(output.nQuantizationValue)).c_str()
                       : "");
            printf("    Memory %s\n", describe_memory_type(output.eMemoryType));
            printf("    Size %u\n", output.nSize);
        }
    }

    int parse_npu_mode_from_joint(const AX_CHAR *data, AX_U32 data_size,
                                  AX_NPU_SDK_EX_HARD_MODE_T *pMode)
    {
        AX_NPU_SDK_EX_MODEL_TYPE_T npu_type;

        auto ret = AX_JOINT_GetJointModelType(data, data_size, &npu_type);
        if (AX_ERR_NPU_JOINT_SUCCESS != ret)
        {
            return -1;
        }

        if (AX_NPU_MODEL_TYPE_DEFUALT == npu_type)
        {
            // fprintf(stdout, "[INFO]: Virtual npu was disabled!\n");
        }
        else if (npu_type == AX_NPU_MODEL_TYPE_1_1_1 || npu_type == AX_NPU_MODEL_TYPE_1_1_2)
        {
            *pMode = AX_NPU_SDK_EX_HARD_MODE_T::AX_NPU_VIRTUAL_1_1;
        }
        else
        {
            return -1;
        }
        return ret;
    }

    int alloc_joint_buffer(const AX_JOINT_IOMETA_T *pMeta, AX_JOINT_IO_BUFFER_T *pBuf, AX_JOINT_ALLOC_BUFFER_STRATEGY_T eStrategy = AX_JOINT_ABST_DEFAULT)
    {
        AX_JOINT_IOMETA_T meta = *pMeta;
        auto ret = AX_JOINT_AllocBuffer(&meta, pBuf, eStrategy);
        if (AX_ERR_NPU_JOINT_SUCCESS != ret)
        {
            ALOGE("Cannot allocate memory.");
            return -1;
        }

        return AX_ERR_NPU_JOINT_SUCCESS;
    }

    int free_joint_buffer(AX_JOINT_IO_BUFFER_T *pBuf)
    {
        auto ret = AX_JOINT_FreeBuffer(pBuf);
        if (AX_ERR_NPU_JOINT_SUCCESS != ret)
        {
            ALOGE("Free allocated memory failed.");
            return -1;
        }
        return AX_ERR_NPU_JOINT_SUCCESS;
    }

    AX_JOINT_IO_BUFFER_T *prepare_io(const size_t &size, AX_JOINT_IO_T &io, const AX_JOINT_IO_INFO_T *io_info,
                                     AX_JOINT_ALLOC_BUFFER_STRATEGY_T eStrategy = AX_JOINT_ABST_DEFAULT,
                                     const uint32_t &batch = 1)
    {
        std::memset(&io, 0, sizeof(io));

        io.nInputSize = io_info->nInputSize;

        if (1 != io.nInputSize)
        {
            ALOGE("Only single input was accepted(got %u).", io.nInputSize);
            return nullptr;
        }

        io.pInputs = new AX_JOINT_IO_BUFFER_T[io.nInputSize];

        // fill input
        const AX_JOINT_IOMETA_T *pMeta = io_info->pInputs;
        AX_JOINT_IO_BUFFER_T *pBuf = io.pInputs;

        if (pMeta->nShapeSize <= 0)
        {
            ALOGE("Dimension(%u) of shape is not allowed.", (uint32_t)pMeta->nShapeSize);
            return nullptr;
        }

        auto actual_data_size = pMeta->nSize / pMeta->pShape[0] * batch;
        if (size != actual_data_size)
        {
            ALOGE("The buffer size is not equal to model input(%s) size(%u vs %u).", io_info->pInputs[0].pName,
                  (uint32_t)size, actual_data_size);
            return nullptr;
        }

        pBuf->nSize = pMeta->nSize;

        // deal with output
        io.nOutputSize = io_info->nOutputSize;
        io.pOutputs = new AX_JOINT_IO_BUFFER_T[io.nOutputSize];
        for (size_t i = 0; i < io.nOutputSize; ++i)
        {
            const AX_JOINT_IOMETA_T *pMeta = io_info->pOutputs + i;
            AX_JOINT_IO_BUFFER_T *pTmpBuf = io.pOutputs + i;
            alloc_joint_buffer(pMeta, pTmpBuf, eStrategy);
        }
        return pBuf;
    }

    int query_model_input_size(const AX_JOINT_IO_INFO_T *io_info, std::array<int, 2> &input_size, AX_U32 &nModelInputSize,
                               AX_NPU_CV_FrameDataType &eDtype)
    {
        int height = 0;
        int width = 0;
        int channel = 0;
        int data_type_size = 0;
        auto &input = io_info->pInputs[0];
        switch (input.eLayout)
        {
        case AX_JOINT_TENSOR_LAYOUT_T::JOINT_TENSOR_LAYOUT_NHWC:
            height = input.pShape[1];
            width = input.pShape[2];
            channel = input.pShape[3];
            nModelInputSize = input.nSize;
            break;
        case AX_JOINT_TENSOR_LAYOUT_T::JOINT_TENSOR_LAYOUT_NCHW:
            channel = input.pShape[1];
            height = input.pShape[2];
            width = input.pShape[3];
            nModelInputSize = input.nSize;
            break;
        default:
            return -1;
        }

        switch (input.eDataType)
        {
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT8:
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT8:
            data_type_size = 1;
            break;
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT16:
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT16:
            data_type_size = 2;
            break;
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_FLOAT32:
            data_type_size = 4;
            break;
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_SINT32:
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT32:
            data_type_size = 4;
            break;
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_FLOAT64:
            data_type_size = 8;
            break;
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT10_PACKED:
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT12_PACKED:
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT14_PACKED:
        case AX_JOINT_DATA_TYPE_T::AX_JOINT_DT_UINT16_PACKED:
        default:
            return -1;
        }

        if (channel == 0 || height == 0 || width == 0 || nModelInputSize == 0 || data_type_size == 0)
        {
            return -1;
        }

        switch (input.pExtraMeta->eColorSpace)
        {
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_BGR:
            input_size[0] = nModelInputSize / (width * data_type_size * 3);
            input_size[1] = width;
            eDtype = AX_NPU_CV_FDT_BGR;
            break;
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RGB:
            input_size[0] = nModelInputSize / (width * data_type_size * 3);
            input_size[1] = width;
            eDtype = AX_NPU_CV_FDT_RGB;
            break;
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_NV12:
            input_size[0] = (nModelInputSize * 2) / (width * data_type_size * 3);
            input_size[1] = width;
            eDtype = AX_NPU_CV_FDT_NV12;
            break;
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_GRAY:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_NV21:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_YUV444:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW8:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW10:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW12:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW14:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_RAW16:
        case AX_JOINT_COLOR_SPACE_T::AX_JOINT_CS_FEATUREMAP:
        default:
            return -1;
        }

        return 0;
    }

    uint32_t get_image_stride_w(const AX_NPU_CV_Image *pImg)
    {
        if (pImg->tStride.nW == 0)
        {
            return pImg->nWidth;
        }
        else if (pImg->tStride.nW >= pImg->nWidth)
        {
            if ((pImg->eDtype == AX_NPU_CV_FDT_NV12 || pImg->eDtype == AX_NPU_CV_FDT_NV21) && pImg->tStride.nW % 2 != 0)
            {
                ALOGE("Invalid param: image stride_w %d should be even for NV12/NV21", pImg->tStride.nW);
            }

            return pImg->tStride.nW;
        }
        else
        {
            ALOGE("Invalid param: image stride_w %d not less than image width %d", pImg->tStride.nW, pImg->nWidth);
        }
        return pImg->nWidth;
    }

    int get_image_data_size(const AX_NPU_CV_Image *img)
    {
        int stride_w = get_image_stride_w(img);
        switch (img->eDtype)
        {
        case AX_NPU_CV_FDT_NV12: // FIXME
        case AX_NPU_CV_FDT_NV21:
            return int(stride_w * img->nHeight * 3 / 2);

        case AX_NPU_CV_FDT_RGB:
        case AX_NPU_CV_FDT_BGR:
        case AX_NPU_CV_FDT_YUV444:
            return int(stride_w * img->nHeight * 3);

        case AX_NPU_CV_FDT_RGBA:
            return int(stride_w * img->nHeight * 4);

        case AX_NPU_CV_FDT_GRAY:
            return int(stride_w * img->nHeight * 1);

        default:
            ALOGE("unsupported color space %d to calculate image data size", (int)img->eDtype);
            return 0;
        }
    }

    AX_NPU_CV_Image *alloc_cv_image(const ax_cv_image &input)
    {
        int ret = 0;
        auto dst_image = new AX_NPU_CV_Image;

        if (!dst_image)
        {
            ALOGE("error AX_NPU_CV_Image");
            return nullptr;
        }

        memset(dst_image, 0x00, sizeof(AX_NPU_CV_Image));
        dst_image->nWidth = input.w;
        dst_image->nHeight = input.h;
        dst_image->tStride.nW = input.stride_w;
        dst_image->eDtype = input.color_space;
        dst_image->nSize = get_image_data_size(dst_image);
        ret = AX_SYS_MemAlloc((AX_U64 *)&dst_image->pPhy, (void **)&dst_image->pVir, dst_image->nSize, 0x100, (AX_S8 *)"SAMPLE-CV");

        if (ret != AX_ERR_NPU_JOINT_SUCCESS)
        {
            ALOGE("error alloc image sys mem %x", ret);
            return nullptr;
        }

        return dst_image;
    }

    int free_cv_image(AX_NPU_CV_Image *image)
    {
        int ret = AX_SYS_MemFree((AX_U64)image->pPhy, (void *)image->pVir);

        delete image;

        if (ret != AX_ERR_NPU_JOINT_SUCCESS)
        {
            ALOGE("[ERR] error free %x", ret);
            return ret;
        }

        return 0;
    }

    int npu_crop_resize(const AX_NPU_CV_Image *input_image, AX_NPU_CV_Image *output_image, AX_NPU_CV_Box *box,
                        AX_NPU_SDK_EX_MODEL_TYPE_T model_type, AX_NPU_CV_ImageResizeAlignParam horizontal,
                        AX_NPU_CV_ImageResizeAlignParam vertical)
    {
        AX_NPU_CV_Color color;
        color.nYUVColorValue[0] = 0;
        color.nYUVColorValue[1] = 128;
        AX_NPU_SDK_EX_MODEL_TYPE_T virtual_npu_mode_type = model_type;

        AX_NPU_CV_Box *ppBox[1];
        ppBox[0] = box;

        int ret = AX_NPU_CV_CropResizeImage(virtual_npu_mode_type, input_image, 1, &output_image, ppBox, horizontal, vertical, color);

        if (ret != AX_NPU_DEV_STATUS_SUCCESS)
        {
            ALOGE("AX_NPU_CV_CropResizeImage err code: %X", ret);
            return ret;
        }

        return 0;
    }

    int npu_convert_color_space(const AX_NPU_CV_Image *input_image, AX_NPU_CV_Image *output_image, AX_NPU_SDK_EX_MODEL_TYPE_T model_type)
    {
        AX_NPU_SDK_EX_MODEL_TYPE_T virtual_npu_mode_type = model_type;

        int ret = AX_NPU_CV_CSC(virtual_npu_mode_type, input_image, output_image);

        if (ret != AX_NPU_DEV_STATUS_SUCCESS)
        {
            ALOGE("AX_NPU_CV_convert_color_space err code: %X", ret);
            return ret;
        }
        return 0;
    }

    int FrameCropResize(const AX_VIDEO_FRAME_S *pstFrame, AX_NPU_CV_Image *pDstNV12)
    {
        auto ConvFrame = [](const AX_VIDEO_FRAME_S *pSrc) -> AX_NPU_CV_Image
        {
            AX_NPU_CV_Image t = {0};
            t.nWidth = pSrc->u32Width;
            t.nHeight = pSrc->u32Height;
            t.eDtype = AX_NPU_CV_FDT_NV12;
            t.tStride.nW = (0 == pSrc->u32PicStride[0]) ? t.nWidth : pSrc->u32PicStride[0];
            t.nSize = pSrc->u32FrameSize; // t.tStride.nW * t.nHeight * 3 / 2;
            t.pPhy = pSrc->u64PhyAddr[0];
            t.pVir = (AX_U8 *)pSrc->u64VirAddr[0];
            return t;
        };

        AX_NPU_CV_Image tSrcFrame = ConvFrame(pstFrame);
        npu_crop_resize(&tSrcFrame, pDstNV12, NULL, AX_NPU_MODEL_TYPE_1_1_1, AX_NPU_CV_IMAGE_HORIZONTAL_CENTER, AX_NPU_CV_IMAGE_VERTICAL_CENTER);
        return 0;
    }

    int copy_to_device(AX_NPU_CV_Image *pImage, AX_JOINT_IO_BUFFER_T *pBuf)
    {
        pBuf->phyAddr = (AX_ADDR)pImage->pPhy;
        pBuf->pVirAddr = (AX_VOID *)pImage->pVir;
        pBuf->nSize = (AX_U32)pImage->nSize;

        return AX_ERR_NPU_JOINT_SUCCESS;
    }

    int cache_io_flush(const AX_JOINT_IO_BUFFER_T *pBuf)
    {
        if (pBuf->phyAddr != 0)
        {
            AX_SYS_MemFlushCache(pBuf->phyAddr, pBuf->pVirAddr, pBuf->nSize);
        }

        return 0;
    }
    int RunBodyDetection(SampleRunJoint_PRIV_CTX_S &stCtx, AX_NPU_CV_Image *pImage)
    {
        int nRet = 0;
        // AX_JOINT_IO_SETTING_T joint_io_setting;
        AX_JOINT_COMPONENT_T *joint_comps;
        AX_U32 joint_comp_size = 0;

        // 3.1. fill input & prepare to inference
        auto ret = copy_to_device(pImage, stCtx.pstIOBuffer);

        if (AX_ERR_NPU_JOINT_SUCCESS != ret)
        {
            ALOGE("copy_to_device failed.");
            nRet = -1;
            goto EXIT;
        }

        // 3.2. run & benchmark
        ret = AX_JOINT_RunSync(stCtx.pHandle, stCtx.stJointCtx, &stCtx.stJointIoArr);

        if (AX_ERR_NPU_JOINT_SUCCESS != ret)
        {
            ALOGE("body Inference failed(%d).", ret);
            nRet = -1;
            goto EXIT;
        }

        ret = AX_JOINT_ADV_GetComponents(stCtx.stJointCtx, &joint_comps, &joint_comp_size);
        if (AX_ERR_NPU_JOINT_SUCCESS != ret)
        {
            ALOGE("body Get components after run failed.");
            nRet = -1;
            goto EXIT;
        }

    EXIT:
        return nRet;
    }

    int BodyPostProcess(SampleRunJoint_PRIV_CTX_S &stCtx, AX_U32 nSrcHeight, AX_U32 nSrcWidth,
                        std::vector<detection::Object> &object_bbox, AX_F32 fConfidence)
    {
        // post process
        // 4.1 get bbox
        std::vector<detection::Object> proporsal;
        std::vector<std::vector<int>> stride = {{8}, {16}, {32}};
        // auto io_info = AX_JOINT_GetIOInfo(stCtx.pHandle);

        for (int i = 0; i < 3; ++i)
        {
            auto &info = stCtx.stJointIoArr.pOutputs[2 * i];
            auto &info_2 = stCtx.stJointIoArr.pOutputs[2 * i + 1];
            auto feat_ptr = (float *)info.pVirAddr;
            auto feat_kps_ptr = (float *)info_2.pVirAddr;
            int32_t stride = (1 << i) * 8;
            detection::generate_proposals_yolov8_pose_native(stride, feat_ptr, feat_kps_ptr, PROB_THRESHOLD, proporsal, 640, 640, 15, 2);
        }

        detection::get_out_bbox_kps(proporsal, object_bbox, NMS_THRESHOLD, stCtx.arrInputSize[0], stCtx.arrInputSize[1], nSrcHeight,
                                    nSrcWidth);
        return 0;
    }

    int BodySegProcess(SampleRunJoint_PRIV_CTX_S &stCtx, AX_U32 nSrcHeight, AX_U32 nSrcWidth,
                       std::vector<detection::Object> &object_bbox, AX_F32 fConfidence)
    {
        std::vector<detection::Object> proporsal;
        std::vector<std::vector<int>> stride = {{8}, {16}, {32}};

        float *output_ptr[3] = {(float *)stCtx.stJointIoArr.pOutputs[0].pVirAddr,      // 1*80*80*144
                                (float *)stCtx.stJointIoArr.pOutputs[1].pVirAddr,      // 1*40*40*144
                                (float *)stCtx.stJointIoArr.pOutputs[2].pVirAddr};     // 1*20*20*144
        float *output_seg_ptr[3] = {(float *)stCtx.stJointIoArr.pOutputs[3].pVirAddr,  // 1*80*80*32
                                    (float *)stCtx.stJointIoArr.pOutputs[4].pVirAddr,  // 1*40*40*32
                                    (float *)stCtx.stJointIoArr.pOutputs[5].pVirAddr}; // 1*20*20*32
        ALOGE("BodySegProcess output_seg_ptr DONE");
        for (int i = 0; i < 3; ++i)
        {
            auto feat_ptr = output_ptr[i];
            auto feat_seg_ptr = output_seg_ptr[i];
            int32_t stride = (1 << i) * 8;
            detection::generate_proposals_yolov8_seg_native(stride, feat_ptr, feat_seg_ptr, PROB_THRESHOLD, proporsal, 640, 640, 1);
        }
        for (const auto &proposal : proporsal)
        {
            ALOGE("Rect: (%f, %f, %f, %f)", proposal.rect.x, proposal.rect.y, proposal.rect.width, proposal.rect.height);
            ALOGE("Label: %d", proposal.label);
            ALOGE("Probability: %f", proposal.prob);

            // Print mask size
            ALOGE("Mask size: (%d, %d)", proposal.mask.rows, proposal.mask.cols);

            // Print mask_feat size
            ALOGE("Mask Feature size: %zu", proposal.mask_feat.size());

            // Print kps_feat size
            ALOGE("Keypoints Feature size: %zu", proposal.kps_feat.size());

            // Separator line
            ALOGE("----------------------");
        }

        ALOGE("BodySegProcess generate_proposals_yolov8_seg_native done");
        ALOGE("BodySegProcess proporsal size %zu", proporsal.size());

        // 1*32*160*160
        auto mask_proto_ptr = (float *)stCtx.stJointIoArr.pOutputs[6].pVirAddr;
        //  int mask_proto_dim, int mask_stride, const float nms_threshold, int letterbox_rows, int letterbox_cols, int src_rows, int src_cols)
        detection::get_out_bbox_mask(proporsal, object_bbox, mask_proto_ptr, 32, 4, NMS_THRESHOLD, 640, 640, nSrcHeight, nSrcWidth);

        // detection::get_out_bbox_mask(proporsal, object_bbox, 8, mask_proto_ptr, 32, 4, NMS_THRESHOLD,
        //                          640, 640,nSrcHeight, nSrcWidth);

        std::sort(object_bbox.begin(), object_bbox.end(),
                  [&](detection::Object &a, detection::Object &b)
                  {
                      return a.rect.area() > b.rect.area();
                  });
    }

}

int SampleRunJointInit(const AX_CHAR *pstrModel, AX_VOID **ppJointHandle, const AX_S32 s32ModelType)
{
    if (s32ModelType != 1 && s32ModelType != 2)
    {
        ALOGE("s32ModelType error %d", s32ModelType);
        return -1;
    }
    g_stInstCtx.modelType = s32ModelType;
    if (!pstrModel)
    {
        ALOGE("models file input error");
        return -1;
    }

    if (!ppJointHandle)
    {
        ALOGE("invalid param");
        return -1;
    }
    // 1.1 read model file to buffer
    AX_JOINT_HANDLE pJointHandle = nullptr;
    std::vector<char> model_buffer;

    if (!read_file(pstrModel, model_buffer))
    {
        ALOGE("read file %s failed", pstrModel);
        return -1;
    }

    // 1.2 parse model from buffer
    AX_JOINT_SDK_ATTR_T attr;
    memset(&attr, 0, sizeof(AX_JOINT_SDK_ATTR_T));

    auto ret = parse_npu_mode_from_joint(model_buffer.data(), model_buffer.size(), &attr.eNpuMode);

    if (AX_ERR_NPU_JOINT_SUCCESS != ret)
    {
        ALOGE("parse model(%s) fail", pstrModel);
        return -1;
    }

    // 1.3 init model
    ret = AX_JOINT_Adv_Init(&attr);
    if (AX_ERR_NPU_JOINT_SUCCESS != ret)
    {
        ALOGE("init model(%s) fail", pstrModel);
        return -1;
    }

    auto deinit_joint = [&pJointHandle]()
    {
        AX_JOINT_DestroyHandle(pJointHandle);
        AX_JOINT_Adv_Deinit();
        return -1;
    };

    // 1.4 the real init processing
    ret = AX_JOINT_CreateHandle(&pJointHandle, model_buffer.data(), model_buffer.size());

    if (AX_ERR_NPU_JOINT_SUCCESS != ret)
    {
        ALOGE("Create model(%s) Run-Joint handler fail", pstrModel);

        // drop the model buffer
        std::vector<char>().swap(model_buffer);

        return deinit_joint();
    }

    // 1.5 get the version of toolkit (optional)
    const AX_CHAR *version = AX_JOINT_GetModelToolsVersion(pJointHandle);
    ALOGN("model tools version: %s", version);

    // 1.6 drop the model buffer
    std::vector<char>().swap(model_buffer);

    // 1.7 show model io brief info (optional)
    auto io_info = AX_JOINT_GetIOInfo(pJointHandle);
    brief_io_info(io_info);

    // 1.8 query model io info
    AX_U32 nModelInputSize = 0;
    std::array<int, 2> input_size;
    AX_NPU_CV_FrameDataType eDtype = AX_NPU_CV_FDT_NV12;
    ret = query_model_input_size(io_info, input_size, nModelInputSize, eDtype);
    if (0 != ret)
    {
        ALOGE("model(%s) query model input size fail", pstrModel);
        return deinit_joint();
    }

    // check model data type
    if (eDtype != AX_NPU_CV_FDT_NV12)
    {
        ALOGE("model(%s) data type is: %d, not NV12", pstrModel, eDtype);
        return deinit_joint();
    }

    // 1.9 prepare io
    g_stInstCtx.pstIOBuffer = prepare_io(nModelInputSize, g_stInstCtx.stJointIoArr, io_info, AX_JOINT_ABST_CACHED);

    if (!g_stInstCtx.pstIOBuffer)
    {
        ALOGE("model(%s) prepare io fail", pstrModel);
        return deinit_joint();
    }

    if (!g_stInstCtx.pDstNV12)
    {
        ax_cv_image dstImage;
        dstImage.h = input_size[0];
        dstImage.w = input_size[1];
        dstImage.stride_w = dstImage.w;
        dstImage.color_space = AX_NPU_CV_FDT_NV12;
        g_stInstCtx.pDstNV12 = alloc_cv_image(dstImage);
    }

    if (!g_stInstCtx.pDstNV12)
    {
        ALOGE("model(%s) prepare dst buffer fail", pstrModel);
        return deinit_joint();
    }

    // 1.10 create context
    AX_JOINT_EXECUTION_CONTEXT_SETTING_T joint_ctx_settings;
    memset(&g_stInstCtx.stJointCtx, 0, sizeof(g_stInstCtx.stJointCtx));
    memset(&joint_ctx_settings, 0, sizeof(joint_ctx_settings));
    ret = AX_JOINT_CreateExecutionContextV2(pJointHandle, &g_stInstCtx.stJointCtx, &joint_ctx_settings);

    if (AX_ERR_NPU_JOINT_SUCCESS != ret)
    {
        ALOGE("model(%s) Create execution context fail", pstrModel);
        return deinit_joint();
    }

    // 1.11 get the init profile info.
    AX_JOINT_COMPONENT_T *joint_comps;
    AX_U32 joint_comp_size = 0;
    ret = AX_JOINT_ADV_GetComponents(g_stInstCtx.stJointCtx, &joint_comps, &joint_comp_size);

    if (AX_ERR_NPU_JOINT_SUCCESS != ret)
    {
        ALOGE("Get components failed.");
        AX_JOINT_DestroyExecutionContext(g_stInstCtx.stJointCtx);
        return deinit_joint();
    }

    // 1.12 init complete
    *ppJointHandle = pJointHandle;
    g_stInstCtx.pHandle = pJointHandle;
    g_stInstCtx.arrInputSize = input_size;
    g_stInstCtx.eDtype = eDtype;
    g_stInstCtx.stJointIoArr.pIoSetting = &g_stInstCtx.stJointIoSetting;
    memset(&g_stInstCtx.stJointIoSetting, 0x00, sizeof(g_stInstCtx.stJointIoSetting));

    return 0;
}

int SampleRunJointDeInit(AX_VOID *joint_handle)
{
    auto DestroyJoint = [&](SampleRunJoint_PRIV_CTX_S &stCtx)
    {
        if (stCtx.stJointIoArr.pInputs)
        {
            delete[] stCtx.stJointIoArr.pInputs;
        }

        if (stCtx.stJointIoArr.pOutputs)
        {
            for (size_t i = 0; i < stCtx.stJointIoArr.nOutputSize; ++i)
            {
                AX_JOINT_IO_BUFFER_T *pBuf = stCtx.stJointIoArr.pOutputs + i;
                free_joint_buffer(pBuf);
            }

            delete[] stCtx.stJointIoArr.pOutputs;
        }

        AX_JOINT_DestroyExecutionContext(stCtx.stJointCtx);
        AX_JOINT_DestroyHandle(stCtx.pHandle);
        AX_JOINT_Adv_Deinit();
    };

    if (g_stInstCtx.pHandle == joint_handle)
    {
        DestroyJoint(g_stInstCtx);
        g_stInstCtx.pHandle = nullptr;
    }
    else
    {
        ALOGE("joint_handle: %p invalid", joint_handle);
        return -1;
    }

    if (g_stInstCtx.pDstNV12)
    {
        free_cv_image(g_stInstCtx.pDstNV12);
        g_stInstCtx.pDstNV12 = nullptr;
    }

    return 0;
}

int SampleRunJointProcess(AX_VOID *joint_handle, const AX_VIDEO_FRAME_S *pstFrame, sampleRunJoint_RESULT_S **ppResult)
{

    if (g_stInstCtx.pHandle != joint_handle)
    {
        ALOGE("joint_handle: %p invalid", joint_handle);

        return -1;
    }

    // check memory
    if (pstFrame->u64VirAddr[0] == 0 || pstFrame->u64PhyAddr[0] == 0)
    {
        ALOGW("VirAddr(%llx) && u64PhyAddr(%llx) should be set", pstFrame->u64VirAddr[0], pstFrame->u64PhyAddr[0]);
        return -1;
    }

    // check eImgType
    if (AX_YUV420_SEMIPLANAR != pstFrame->enImgFormat)
    {
        ALOGE("Only support NV12");
        return -1;
    }

    if (!ppResult)
    {
        ALOGE("invalid param");
        return -1;
    }

    // 2. pre process
    FrameCropResize(pstFrame, g_stInstCtx.pDstNV12);
    // 3. run body process
    AX_U32 nWidth = pstFrame->u32Width;
    AX_U32 nHeight = pstFrame->u32Height;

    std::vector<detection::Object> bodyOutObject;
    auto nRet = RunBodyDetection(g_stInstCtx, g_stInstCtx.pDstNV12);

    // std::count << nRet << std::endl;
    // printf("%d\n",nRet);
    // ALOGE("nRet,%d",nRet);

    // 4. post process
    if (0 != nRet)
    {
        ALOGE("Run Body Detection fail");
        return -1;
    }

    if (g_stInstCtx.modelType == 1)
    {
        nRet = BodyPostProcess(g_stInstCtx, nHeight, nWidth, bodyOutObject, PROB_THRESHOLD);

        // 5. result output
        sampleRunJoint_RESULT_S *pstResult = new sampleRunJoint_RESULT_S;

        if (!pstResult)
        {
            ALOGE("memory alloc fail");
            return -1;
        }

        memset(pstResult, 0x00, sizeof(sampleRunJoint_RESULT_S));
        pstResult->bHasLandmark = 1;
        pstResult->bHasMask = 0;
        pstResult->nObjectSize = bodyOutObject.size();

        AX_U32 nObjectIndex = 0;

        if (pstResult->nObjectSize > 0)
        {
            pstResult->pstObjectItems = new sampleRunJoint_OBJECT_ITEM_S[pstResult->nObjectSize];
            if (!pstResult->pstObjectItems)
            {
                pstResult->nObjectSize = 0;
                ALOGW("alloc Object Item fail");
                nRet = -1;
                goto EXIT;
            }
            memset(pstResult->pstObjectItems, 0x00, sizeof(sampleRunJoint_OBJECT_ITEM_S) * pstResult->nObjectSize);
        }

        for (size_t i = 0; i < pstResult->nObjectSize; i++)
        {
            const detection::Object &obj = bodyOutObject[i];

            pstResult->pstObjectItems[nObjectIndex].stRect.x = (float)obj.rect.x;
            pstResult->pstObjectItems[nObjectIndex].stRect.y = (float)obj.rect.y;
            pstResult->pstObjectItems[nObjectIndex].stRect.w = (float)obj.rect.width;
            pstResult->pstObjectItems[nObjectIndex].stRect.h = (float)obj.rect.height;
            pstResult->pstObjectItems[nObjectIndex].prob = (float)obj.prob;
            pstResult->pstObjectItems[nObjectIndex].nLandmark = (int)obj.kps_feat.size() / 3;
            pstResult->pstObjectItems[nObjectIndex].labelId = (float)obj.label;
            pstResult->pstObjectItems[nObjectIndex].landmark = new sampleRunJoint_POINT_S[pstResult->pstObjectItems[nObjectIndex].nLandmark];
            memset(pstResult->pstObjectItems[nObjectIndex].landmark, 0x00, sizeof(sampleRunJoint_POINT_S) * pstResult->pstObjectItems[nObjectIndex].nLandmark);
            for (size_t j = 0; j < pstResult->pstObjectItems[nObjectIndex].nLandmark; j++)
            {
                pstResult->pstObjectItems[nObjectIndex].landmark[j].x = obj.kps_feat[3 * j];
                pstResult->pstObjectItems[nObjectIndex].landmark[j].y = obj.kps_feat[3 * j + 1];
                pstResult->pstObjectItems[nObjectIndex].landmark[j].score = obj.kps_feat[3 * j + 2];
            }
            if (obj.label < (int)KEYPOINT_CLASS_NAMES.size())
            {
                strcpy(pstResult->pstObjectItems[nObjectIndex].labelName, KEYPOINT_CLASS_NAMES[obj.label].c_str());
            }
            else
            {
                strcpy(pstResult->pstObjectItems[nObjectIndex].labelName, "unknown");
            }
            nObjectIndex++;
        }

        pstResult->nWidth = pstFrame->u32Width;
        pstResult->nHeight = pstFrame->u32Height;
        pstResult->nObjectSize = nObjectIndex;

        *ppResult = pstResult;
    }
    else if (g_stInstCtx.modelType == 2)
    {
        ALOGE("BodySegProcess Start");
        nRet = BodySegProcess(g_stInstCtx, nHeight, nWidth, bodyOutObject, PROB_THRESHOLD);
        ALOGE("BodySegProcess DONE");

        // 5. result output
        sampleRunJoint_RESULT_S *pstResult = new sampleRunJoint_RESULT_S;

        if (!pstResult)
        {
            ALOGE("memory alloc fail");
            return -1;
        }

        memset(pstResult, 0x00, sizeof(sampleRunJoint_RESULT_S));
        pstResult->bHasLandmark = 0;
        pstResult->bHasMask = 1;
        pstResult->nObjectSize = bodyOutObject.size();

        AX_U32 nObjectIndex = 0;

        if (pstResult->nObjectSize > 0)
        {
            pstResult->pstObjectItems = new sampleRunJoint_OBJECT_ITEM_S[pstResult->nObjectSize];
            if (!pstResult->pstObjectItems)
            {
                pstResult->nObjectSize = 0;
                ALOGW("alloc Object Item fail");
                nRet = -1;
                goto EXIT;
            }
            memset(pstResult->pstObjectItems, 0x00, sizeof(sampleRunJoint_OBJECT_ITEM_S) * pstResult->nObjectSize);
        }
        for (size_t i = 0; i < pstResult->nObjectSize; i++)
        {
            const detection::Object &obj = bodyOutObject[i];

            pstResult->pstObjectItems[nObjectIndex].stRect.x = (float)obj.rect.x;
            pstResult->pstObjectItems[nObjectIndex].stRect.y = (float)obj.rect.y;
            pstResult->pstObjectItems[nObjectIndex].stRect.w = (float)obj.rect.width;
            pstResult->pstObjectItems[nObjectIndex].stRect.h = (float)obj.rect.height;
            pstResult->pstObjectItems[nObjectIndex].prob = (float)obj.prob;
            pstResult->pstObjectItems[nObjectIndex].labelId = (float)obj.label;
            pstResult->pstObjectItems[nObjectIndex].mMask.data = obj.mask.data;
            pstResult->pstObjectItems[nObjectIndex].mMask.w = obj.mask.cols;
            pstResult->pstObjectItems[nObjectIndex].mMask.h = obj.mask.rows;
            pstResult->pstObjectItems[nObjectIndex].mMask.c = obj.mask.channels();
            pstResult->pstObjectItems[nObjectIndex].mMask.s = obj.mask.step1();
            if (obj.label < (int)SEG_CLASS_NAMES.size())
            {
                strcpy(pstResult->pstObjectItems[nObjectIndex].labelName, SEG_CLASS_NAMES[obj.label].c_str());
            }
            else
            {
                strcpy(pstResult->pstObjectItems[nObjectIndex].labelName, "unknown");
            }
            nObjectIndex++;
        }

        pstResult->nWidth = pstFrame->u32Width;
        pstResult->nHeight = pstFrame->u32Height;
        pstResult->nObjectSize = nObjectIndex;

        *ppResult = pstResult;
    }
EXIT:
    return nRet;
}

int SampleRunJointReleaseResult(AX_VOID *pResult)
{
    if (pResult)
    {
        sampleRunJoint_RESULT_S *pstResult = (sampleRunJoint_RESULT_S *)pResult;

        if (pstResult->pstObjectItems)
        {
            delete[] pstResult->pstObjectItems;
        }

        delete pstResult;
    }

    return 0;
}
