#include "reid.h"

#include <stdio.h>
#include <stdexcept>
// #include <cuda_runtime_api.h>
#include <torch/torch.h>
#include "opencv2/opencv.hpp"

reid::reid(const int gpu_index, const char *module_path)
:is_gpu_(gpu_index >= 0)
,start_(false)
{
    if(init(gpu_index, module_path) < 0){
        throw std::runtime_error("init swscale error");
    }
}

reid::~reid()
{
    release();
}

int reid::init(const int gpu_index, const char *module_path){
    const int gpu_id = gpu_index;
    if(gpu_id == -1){
        module_ = torch::jit::load(module_path);
        module_.to(torch::kCPU);
        module_.eval();
        is_gpu_ = false;
        printf("==========load CPU\n");
    }else if(torch::cuda::is_available() && torch::cuda::device_count() >= gpu_id){
        // cudaSetDevice(gpu_id);
        module_ = torch::jit::load(module_path, torch::Device(torch::DeviceType::CUDA, gpu_id));
        module_.to(torch::kCUDA);
        module_.eval();
        is_gpu_ = true;
        printf("==========load GPU\n");
    }else{
        printf("reid use gpu %d, error\n", gpu_id);
        return -1;
    }
    start_ = true;
    return 0;
}

void reid::release(){

}

int reid::extract(unsigned char *img, const int w, const int h, const int c, float *feature){
    if (!start_ || !img) return -1;

    printf("&&&&&& Start ReID Image Size: %dx%d\n", w, h);
    printf("\n\n\n\n=============================================");

    cv::Mat matImg(h, w, CV_8UC3, img);

    //cv::VideoCapture capture2;
    //cv::Mat matImg;
    //matImg = capture2.open("rtsp://admin:a1234567@192.168.5.34:554/h264/ch1/main/av_stream");
    //capture2.read(matImg);
    


    // scale image to fit
    cv::Size scale(256, 128);
    cv::resize(matImg, matImg, scale);

    // vector<int> compression_params;
    // compression_params.push_back(cv::IMWRITE_PNG_COMPRESSION);
    // compression_params.push_back(9);

    //cv::imwrite("0000.png", matImg, compression_params);
    matImg.convertTo(matImg, CV_32FC3, 1.0f/255.0f);
    printf("&&&&&& matImg.convertTo\n");

    auto input_tensor = torch::from_blob(matImg.data, {1, 256, 128, 3});
    printf("&&&&&& torch::from_blob\n");

    input_tensor = input_tensor.permute({0, 3, 1, 2});
    printf("&&&&&& input_tensor.permute\n");

    input_tensor[0][0] = input_tensor[0][0].sub_(0.485).div_(0.229);
    printf("&&&&&& input_tensor[0][0]\n");

    input_tensor[0][1] = input_tensor[0][1].sub_(0.456).div_(0.224);
    printf("&&&&&& input_tensor[0][1]\n");

    input_tensor[0][2] = input_tensor[0][2].sub_(0.406).div_(0.225);
    printf("&&&&&& input_tensor[0][2]\n");

    if(is_gpu_)
    {
        input_tensor = input_tensor.to(at::kCUDA);
        printf("&&&&&& kCUDA\n");
    }else
    {
        printf("------ kCPU)\n");
    }


    torch::Tensor human_feat = module_.forward({input_tensor}).toTensor();
    printf("------human_feat---\n");
    torch::Tensor query_feat;

    if(is_gpu_){        
        printf("------is_gpu_--\n");
        query_feat = human_feat.cpu();
    }
    else
        query_feat = human_feat;
 

    printf("&&&&&& query_feat\n");

    auto foo_one = query_feat.accessor<float,2>();
    printf("&&&&&& query_feat.accessor\n");

    float f_size = -0.727412;
    for (int64_t i = 0; i < foo_one.size(0); i++) {
        printf("------foo_one.size(0)::%d\n",foo_one.size(0));
        auto a1 = foo_one[i];
        for (int64_t j = 0; j < a1.size(0); j++) {
            feature[j] = a1[j];
        }
    }

    // for (int k = 0; k < 3; ++k) {
    //     printf("--extractor---human_feats------%f",feature[k+2000]);
    // }

    printf("\n\n\n\n\n----------------------------------------------------------------------------\n");

    printf("&&&&&& End ReID With CUDA %d\n", is_gpu_);

    return 0;
}

float reid::compare(const float *feat1, const float *feat2){
    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] = feat1[i];
        gallery_feat[0][i] = feat2[i];
    }
 
    if(is_gpu_) {
        query_feat = query_feat.cuda();
        gallery_feat = gallery_feat.cuda();
    }
  
    torch::Tensor a_similarity = torch::cosine_similarity(query_feat, gallery_feat);
 
    if(is_gpu_)
        a_similarity = a_similarity.cpu();
 
    auto foo_one = a_similarity.accessor<float,1>();
 
    float f_distance = foo_one[0];
 
    return f_distance;
}