//
// Created by Scheaven on 2020/1/3.
//

#include "reid_feature.h"
#include <cuda_runtime_api.h>
#include <torch/torch.h>

bool ReID_Feature::ReID_init(int gpu_id)
{
    if(gpu_id == -1){
        this->module = torch::jit::load(MODEL_PATH);
        this->module.to(torch::kCPU);
        this->module.eval();
        this->is_gpu = false;
    }else if(torch::cuda::is_available() && torch::cuda::device_count() >= gpu_id)
    {
        cudaSetDevice(gpu_id);
        cout << "model loading::" << HUMAN_FEATS << endl;
        this->module = torch::jit::load(MODEL_PATH, torch::Device(torch::DeviceType::CUDA,gpu_id));
        this->module.to(torch::kCUDA);
        this->module.eval();
        this->is_gpu = true;
    }else{
        return false;
    }
    return true;
}

int ReID_Feature::ReID_size()
{
    int size = 2048;
    return size;

}

bool ReID_Feature::ReID_extractor(unsigned char *pBuf, float *pFeature)
{
    auto input_tensor = torch::from_blob(pBuf, {1, 256, 128, 3});
    input_tensor = input_tensor.permute({0, 3, 1, 2});
    input_tensor[0][0] = input_tensor[0][0].sub_(0.485).div_(0.229);
    input_tensor[0][1] = input_tensor[0][1].sub_(0.456).div_(0.224);
    input_tensor[0][2] = input_tensor[0][2].sub_(0.406).div_(0.225);
    if(this->is_gpu)
        input_tensor = input_tensor.to(at::kCUDA);
    torch::Tensor human_feat =this->module.forward({input_tensor}).toTensor();

//    for (int k = 0; k < 20; ++k) {
//        cout << "--extractor---human_feats------" <<human_feat[0][k+2000]<< endl;
//    }
    torch::Tensor query_feat;
    if(this->is_gpu)
        query_feat = human_feat.cpu();
    else
        query_feat = human_feat;

    auto foo_one = query_feat.accessor<float,2>();

    ReID_Utils RET;

    float f_size = -0.727412;
    for (int64_t i = 0; i < foo_one.size(0); i++) {
        auto a1 = foo_one[i];
        for (int64_t j = 0; j < a1.size(0); j++) {
            pFeature[j] = a1[j];
        }
    }

//    cout << "---- end 11-------" << pFeature[0] << endl;
    return true;
}

float ReID_Feature::ReID_Compare(float *pFeature1, float *pFeature2)
{
    torch::Tensor query_feat = torch::zeros({1,2048});
    torch::Tensor gallery_feat = torch::zeros({1,2048});

    for (int i = 0; i < 2048; i++)
    {
        query_feat[0][i] = pFeature1[i];
        gallery_feat[0][i] = pFeature2[i];
    }

    if(this->is_gpu)
    {
        query_feat = query_feat.cuda();
        gallery_feat = gallery_feat.cuda();
    }

//    cout << "-----------------after-----------" << endl;
//    cout << query_feat<< endl;

//    for (int k = 0; k < 20; ++k) {
//        cout << "-query_feat----1111111111------" <<query_feat[0][k+2000]<< endl;
//    }
//
//    cout << "-----------------asdf-----------" << endl;
////    cout << gallery_feat[0][0]<< endl;
//    for (int k = 0; k < 20; ++k) {
//        cout << "-gallery_feat----22222222------" <<gallery_feat[0][k+2000]<< endl;
//    }

    torch::Tensor a_similarity = torch::cosine_similarity(query_feat, gallery_feat);

    if(this->is_gpu)
        a_similarity = a_similarity.cpu();

    auto foo_one = a_similarity.accessor<float,1>();
//    cout << ":::::::::-" << endl;

    float f_distance = foo_one[0];

    return f_distance;

}

void ReID_Feature::ReID_Release()
{
        prinf("release");
//    this->module = nullptr;//加载模型
//    return true;
}