//
// Created by Scheaven on 2020/4/10.
//

#include "FeatureEncoder.h"
#include <vector>
using namespace tensorflow;
using namespace std;

FeatureEncoder* FeatureEncoder::instance = NULL;

FeatureEncoder* FeatureEncoder::getInstance()
{
    if(instance==NULL)
    {
        instance = new FeatureEncoder();
    }
    return instance;
}

FeatureEncoder::~FeatureEncoder()
{
//    delete instance;
}

void FeatureEncoder::release()
{
    delete FeatureEncoder::instance;
    FeatureEncoder::instance = NULL;
}

FeatureEncoder::FeatureEncoder()
{
    Status status;
    GraphDef graph_def;

    SessionOptions sessOptions;
    P_STATICSTRUCT m_config;

    int gpu_device = gpu::nv_get_suitable_gpu();
    string prefix_device = "/gpu:";
    string tracking_gpu = prefix_device + std::to_string(gpu_device);

    tensorflow::graph::SetDefaultDevice(tracking_gpu, &graph_def);
    sessOptions.config.mutable_gpu_options()->set_allow_growth(true);
    sessOptions.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.3);
    status = NewSession(sessOptions,&session);
    if(!status.ok())
    {
        std::cerr << status.ToString() << std::endl;
    }

    status = ReadBinaryProto(Env::Default(), m_config->tracher_model, &graph_def);

    if (!status.ok()) {
        std::cerr << status.ToString() << std::endl;
    }


    status = session->Create(graph_def);
    if (!status.ok()) {
        std::cerr << status.ToString() << std::endl;
    }
}

bool FeatureEncoder::getRectsFeature(const cv::Mat &img, DETECTIONS &detections)
{
    std::vector<cv::Mat> mats;
    for(DETECTION_ROW& dbox:detections)
    {
        cv::Rect rect = cv::Rect(int(dbox.tlwh(0)),int(dbox.tlwh(1)), int(dbox.tlwh(2)), int(dbox.tlwh(3)));
        rect.x -= (rect.height * 0.5 - rect.width) * 0.5;
        rect.width = rect.height * 0.5;
        rect.x = (rect.x >= 0 ? rect.x : 0);
        rect.y = (rect.y >= 0 ? rect.y : 0);
        rect.width = (rect.x + rect.width <= img.cols? rect.width: (img.cols-rect.x));
        rect.height = (rect.y + rect.height <= img.rows? rect.height:(img.rows - rect.y));
        cv::Mat mattmp = img(rect).clone();
        cv::resize(mattmp, mattmp, cv::Size(64, 128));
        mats.push_back(mattmp);
    }

    int count = mats.size();
    Tensor input_image(DT_UINT8, TensorShape({count, 128, 64, 3}));
    mats_buffer(mats, input_image.flat<uint8>().data());

    std::vector<std::pair<tensorflow::string, Tensor>> feed_dict =
            { {"images" , input_image} };

    Status status =  session->Run(feed_dict,{"features"},{}, &output_tensors);

    if (!status.ok()) {
        std::cerr << status.ToString() << std::endl;
        return false;
    }

    float* tensor_buffer = output_tensors[0].flat<float>().data();
    int i = 0;
    for(DETECTION_ROW& dbox : detections) {
        for(int j = 0; j < FEATURE_DIM; j++) {
            dbox.feature[j] = tensor_buffer[i*FEATURE_DIM+j];
        }
        i++;
    }
    return true;
}

void FeatureEncoder::mats_buffer(const std::vector<cv::Mat> &imgs, uint8 *buf)
{
    int pos = 0;
    for(const cv::Mat& img : imgs) {
        int Lenth = img.rows * img.cols * 3;
        int nr = img.rows;
        int nc = img.cols;
        if(img.isContinuous()) {
            nr = 1;
            nc = Lenth;
        }
        for(int i = 0; i < nr; i++) {
            const uchar* inData = img.ptr<uchar>(i);
            for(int j = 0; j < nc; j++) {
                buf[pos] = *inData++;
                pos++;
            }
        }
    }
}