#include "use_reid.h"

using namespace cv;
using namespace std;
using namespace GUD::ALG;

NNIE_REID::NNIE_REID(const gud_reid_config_t reidConf)
{
    cpu_set_t mask;           //CPU核的集合
    cpu_set_t get;            //获取在集合中的CPU
    //int num= sysconf(_SC_NPROCESSORS_CONF);
    int j = 0;
    CPU_ZERO(&mask);    /* 初始化set集，将set置为空*/
    /*将本进程绑定到CPU0上*/
    CPU_SET(3, &mask);   // 3559 core 3

    if (sched_setaffinity(0, sizeof(mask), &mask) == -1) {
        printf("Set CPU affinity failue, ERROR:%s\n", strerror(errno));
        exit(-1);
    } else {
        printf("Set CPU affinity Sucess, Just run NNIE_REID Class\n");
    }

    gallery_num_ = 0;
    gallery_status_ = 0;
    input_w_ = reidConf.input_w;
    input_h_ = reidConf.input_h;
    input_c_ = reidConf.input_c;
    //conf_ = reidConf;
    setGalleyNum_ = reidConf.gallery_num;
    feat_len_ = reidConf.feat_len;
    sim_thd_  = reidConf.sim_thd;
    reid_net_.loadModel(reidConf.model);

    input_size_ = input_w_ * input_h_ * input_c_;
    imgTmpAlloc(&subImg,input_w_,input_h_,GUIDE_U8C1);
    bufTmpAlloc(&tmpbuf_,feat_len_*2);
    outTnsr[0].data = (float*)tmpbuf_.virAddr;
    outTnsr[1].data = (float*)(tmpbuf_.virAddr + feat_len_);

    Dev_Malloc_ReidData(&query_obj_);
    galleryPools_ = (gud_reid_data_t*)malloc(sizeof(gud_reid_data_t) * setGalleyNum_);
    for (int i = 0; i < setGalleyNum_; ++i)
    {
        Dev_Malloc_ReidData(&galleryPools_[i]);
    }
    printf ("[%d %d %d]\n",input_h_,input_w_,input_c_);
    printf ("[%d %d %f %d]\n",setGalleyNum_,feat_len_,sim_thd_,input_size_);
}

NNIE_REID::~NNIE_REID()
{
    Dev_Free_ReidData(&query_obj_);
    for (int i = 0; i < setGalleyNum_; ++i)
    {
        Dev_Free_ReidData(&galleryPools_[i]);
    }
    imgTmpFree(&subImg);
    bufTmpFree(&tmpbuf_);
}

void NNIE_REID::Dev_Malloc_ReidData(gud_reid_data_t *reid_obj)
{
    reid_obj->has_feat = 0;
    reid_obj->feat_len = feat_len_;
    reid_obj->feat = (float *)malloc(sizeof(float) * feat_len_);
    memset(reid_obj->feat, 0, sizeof(float) * feat_len_);
//    reid_obj->imgdata = (gud_u8 *)malloc(sizeof(gud_u8) * input_size_);
//    memset(reid_obj->imgdata, 0, sizeof(gud_u8) * input_size_);
}

void NNIE_REID::Dev_Free_ReidData(gud_reid_data_t *reid_obj)
{
    if (reid_obj->feat != NULL)
        free(reid_obj->feat);
//    if (reid_obj->imgdata != NULL)
//        free(reid_obj->imgdata);
}

void NNIE_REID::Extract_REID_Featrue(GuideImage *pSrcImg, gud_rect_t tag, gud_reid_data_t *reidData)
{
    cv::Mat xCrop,dst_target,dst;

    int tagW = tag.w;
    int tagH = tag.h;
    int xMin = tag.cx - tag.w /2;
    int yMin = tag.cy - tag.h /2;
    int xMax = xMin + tagW;
    int yMax = yMin + tagH;

    int width = pSrcImg->img.au32Stride[0];
    int height = pSrcImg->img.u32Height;
    int bgrSize = width * height * 3;
    char *pSrcBgr = (char*)HI_MPI_SYS_MmapCache(pSrcImg->img.au64PhyAddr[0],bgrSize);
    cv::Mat img = cv::Mat(height,width, CV_8UC3,pSrcBgr);
    cv::Size imgSize = img.size();
    int leftPad = (int)(fmax(0., -xMin));
    int topPad = (int)(fmax(0., -yMin));
    int rightPad = (int)(fmax(0., xMax - imgSize.width + 1));
    int bottomPad = (int)(fmax(0., yMax - imgSize.height + 1));

    xMin = xMin + leftPad;
    yMin = yMin + topPad;
    xMax = xMax + leftPad;
    yMax = yMax + topPad;

    if (topPad == 0 && bottomPad == 0 && leftPad == 0 && rightPad == 0)
    {
        img(cv::Rect(int(xMin), int(yMin), int(xMax - xMin + 1), int(yMax - yMin + 1))).copyTo(xCrop);
    }
    else
    {
        cv::Scalar avgChans = mean(img);
        cv::copyMakeBorder(img, dst, topPad, bottomPad, leftPad, rightPad, cv::BORDER_CONSTANT, avgChans);
        dst(cv::Rect(int(xMin), int(yMin), int(xMax - xMin + 1), int(yMax - yMin + 1))).copyTo(xCrop);
    }

    //img(cv::Rect(xMin, yMin, tagW, tagH)).copyTo(xCrop);
    //std::string regionImg0 = "/nfsroot/001_3559/results/reid_crop.jpg";
    //cv::imwrite(regionImg0, xCrop);

    cv::cvtColor(xCrop, xCrop, cv::COLOR_BGR2GRAY);
    cv::resize(xCrop, dst_target, cv::Size(input_w_, input_h_), 0, 0, cv::INTER_LINEAR);
    //std::string regionImg = "/nfsroot/001_3559/results/reid_init.jpg";
    //cv::imwrite(regionImg, dst_target);

    uint8_t *pTmpBgr = (uint8_t *)subImg.img.au64VirAddr[0];
    cvImageToTensor(dst_target.data, pTmpBgr, 1, input_w_, input_h_);
    reid_net_.run(&subImg,outTnsr);
    reid_output_ = outTnsr[0];
    //printf ("Extract_REID_Featrue sucess!\n");
    //memcpy(reidData->imgdata, dst_target.data, sizeof(gud_u8) * input_size_);
    memcpy(reidData->feat, reid_output_.data, sizeof(float)*feat_len_);
    reidData->feat_len = feat_len_;
    reidData->has_feat = 1;
    reidData->ai_ot    = tag;

    HI_MPI_SYS_Munmap(pSrcBgr,bgrSize);

}

void NNIE_REID::Init_Gallery_Featrue_Pools(GuideImage *pSrcImg, gud_rect_t tag)
{
    printf ("before gallery_num_[%d] ? setGalleyNum_:[%d] [%d]\n",gallery_num_,setGalleyNum_,gallery_status_);
    if (gallery_num_ < setGalleyNum_)
    {
        Extract_REID_Featrue(pSrcImg, tag, &galleryPools_[gallery_num_]);
        gallery_num_++;
        if (gallery_num_==setGalleyNum_)
            gallery_status_ = 1;
    }else{
        gallery_status_ = 1;
    }
    printf ("after gallery_num_[%d] ? setGalleyNum_:[%d] [%d]\n",gallery_num_,setGalleyNum_,gallery_status_);
}

void NNIE_REID::Clear_Gallery_Featrue_Pools()
{

    for (int i = 0; i < setGalleyNum_; ++i) {
        galleryPools_[i].iou_val  = 0;
        galleryPools_[i].has_feat = 0;
        galleryPools_[i].feat_len = 0;
        //memset(galleryPools_[i].imgdata, 0, sizeof(gud_u8) * input_size_);
        memset(galleryPools_[i].feat, 0, sizeof(float)*feat_len_);
    }
    gallery_num_    = 0;
    gallery_status_ = 0;
    printf ("Clear_Gallery_Featrue_Pools set-gallery_num_[%d] [%d]\n",setGalleyNum_,gallery_status_);
}

void NNIE_REID::Clear_Gallery_SetIndx_Featrue_Pools(int setIdx)
{
    if (setIdx < setGalleyNum_){
        for (int i = setIdx; i < setGalleyNum_; ++i) {
            galleryPools_[i].iou_val  = 0;
            galleryPools_[i].has_feat = 0;
            galleryPools_[i].feat_len = 0;
            memset(galleryPools_[i].feat, 0, sizeof(float)*feat_len_);
        }
        gallery_num_    = setIdx;
        gallery_status_ = 0;
        printf ("Clear_Gallery_Featrue_Pools From:[%d-->%d] status:[%d]\n",setIdx,setGalleyNum_,gallery_status_);
    }
}

void NNIE_REID::Compare_QueryToGallery(GuideImage *pSrcImg, gud_rect_t tag, gud_reid_result_t *reidRes)
{
    Extract_REID_Featrue(pSrcImg, tag, &query_obj_);
    float SimAvg = 0;
    float maxSim = 0;
    float minSim = 1000;
    int okNum    = 0;
    //printf (" showTensor query_obj_.feat\n");
    //showTensor(query_obj_.feat,feat_len_,feat_len_);

    for (int i = 0; i < gallery_num_; ++i) {
        //printf (" showTensor galleryPools_[%d] feat\n",i);
        //showTensor(galleryPools_[i].feat,feat_len_,feat_len_);
        float sim = FloatFeatCompare(query_obj_.feat, galleryPools_[i].feat);
        if (sim > sim_thd_)
            okNum++;
        if (sim > maxSim)
            maxSim = sim;
        if (sim < minSim)
            minSim = sim;
        SimAvg+= sim;
        //printf ("query & Gallery[%d] sim:[%.3f]\n",i,sim);
    }
    SimAvg /= gallery_num_;
    reidRes->max_sim = maxSim;
    reidRes->min_sim = minSim;
    reidRes->avg_sim = SimAvg;

    //if (maxSim > sim_thd_ && minSim > sim_thd_ && SimAvg > sim_thd_)
    if (SimAvg > sim_thd_)
        reidRes->match_status = 1;
    else
        reidRes->match_status = 0;
}


float NNIE_REID::FloatFeatCompare(float *feat0, float *feat1)
{
    float cos_sim = 0;
    int n = feat_len_;
    float tmp = 0.0;  //内积
    for (int i = 0; i < n; ++i)
    {
        tmp += feat0[i] * feat1[i];
    }
    float mold = GetMold(feat0, feat_len_) * GetMold(feat1, feat_len_);
    if (mold!=0)
        cos_sim = tmp / mold;
    if (tmp == 0 or mold == 0)
        return 0;
    return cos_sim;
}

float NNIE_REID::GetMold(float* vec, int feat_len)
{
    //求向量的模长
    int n = feat_len;
    double sum = 0.0;
    for (int i = 0; i<n; ++i)
        sum += vec[i] * vec[i];
    return sqrtf(sum);
}

void NNIE_REID::cvImageToTensor(const unsigned char* img_data, unsigned char* tensor, int channels, int height, int width)
{
    int align = 2;
    int dstStride = (width + (align - width % align) % align);

    int hs = height * dstStride;
    int iwc = 0;
    int is = 0;
    int jc = 0;
    for (int i=0; i<height; i++) {
        iwc = i * width * channels;
        is = i * dstStride;
        for (int j=0; j<width; j++) {
            jc = j * channels;
            for (int k=0; k<channels; k++) {
                const size_t offsetCv = iwc + jc + k;
                const size_t offset = k * hs + is + j;
                tensor[offset] = img_data[offsetCv];
            }
        }
    }
}


void NNIE_REID::showTensor(float *TensorIn, int len, int diff)
{
    int lenN = 1;
    printf ("\n***** START ******\n");
    for (int i = 0; i < len; ++i)
    {
        //printf ("%.4f ",Tdata[i]);
        printf ("%.3f ",TensorIn[i]);
        if (i%(diff) == 0 and i!=0){
            printf("\n#[%d]#\n",lenN);
            lenN++;
        }
    }
    printf ("\n***** END ******\n");
}