#include "yolo.h"

#include <pthread.h>

Yolo::Yolo(void) 
{

}

Yolo::~Yolo() 
{
    
}

bool Yolo::Init(char *weight)
{
    bool bret = false;

    model = core.read_model(weight);
    compiled_model = core.compile_model(model, "CPU");
    infer_request = compiled_model.create_infer_request();

    input_tensor = infer_request.get_input_tensor();
    tensor_shape = input_tensor.get_shape();

    inChannel = tensor_shape[1];
    inHeight = tensor_shape[2];
    inWidth = tensor_shape[3];

    bret = true;

    pthread_mutex_init(&yoloMutex, NULL);

    return bret;
}

std::vector<YoloBox> Yolo::Detect(cv::Mat &img, float thresh)
{
    std::vector<YoloBox> yoloBoxes;

    // printf("%d, %d\r\n", inWidth, inHeight);
    pthread_mutex_lock(&yoloMutex);

    std::vector<float> paddings(3);       //scale, half_h, half_w
    std::vector<int> new_shape = {inWidth, inHeight};
    cv::Mat resized_img = Letterbox(img, paddings, new_shape); //resize to (640,640) by Letterbox
    // BGR->RGB, u8(0-255)->f32(0.0-1.0), HWC->NCHW
    cv::Mat blob = cv::dnn::blobFromImage(resized_img, 1 / 255.0, cv::Size(inWidth, inHeight), cv::Scalar(0, 0, 0), true);

    // -------- Step 5. Feed the blob into the input node of YOLOv5 -------
    // Get input port for model with one input
    auto input_port = compiled_model.input();

    // Create tensor from external memory
    ov::Tensor input_tensor(input_port.get_element_type(), input_port.get_shape(), blob.ptr(0));
    // Set input tensor for model with one input
    infer_request.set_input_tensor(input_tensor);

    infer_request.infer();

    auto output_tensor = infer_request.get_output_tensor();
    float* detection = (float*)output_tensor.data();

    auto output_shape = output_tensor.get_shape();

    int rows = output_shape[1];
    //1, 9, 3024
    // printf("%d, %d, %d\r\n", output_shape[0], output_shape[1], output_shape[2]);

    float conf_threshold = thresh;
    float nms_threshold = 0.5;
    std::vector<cv::Rect> rectBoxes;
    std::vector<int> class_ids;
    std::vector<float> class_scores;
    std::vector<float> confidences;

    for(int i=0;i<output_shape[2];i++)
    {
        int maxClassIndex = -1;
        float maxScore = 0.0;

        for(int j=4;j<output_shape[1];j++)
        {
            if(*(detection+j*output_shape[2]+i) > maxScore)
            {
                maxScore = *(detection+j*output_shape[2]+i);
                maxClassIndex = j;
            }
        }

        if(maxScore >= thresh)
        {
            int id = maxClassIndex - 4;

            float cx = *(detection+0*output_shape[2]+i);
            float cy = *(detection+1*output_shape[2]+i);
            float w = *(detection+2*output_shape[2]+i);
            float h = *(detection+3*output_shape[2]+i);
            int left = static_cast<int>((cx - 0.5 * w - paddings[2]) / paddings[0]);
            int top = static_cast<int>((cy - 0.5 * h - paddings[1]) / paddings[0]);
            int width = static_cast<int>(w / paddings[0]);
            int height = static_cast<int>(h / paddings[0]);

            cv::Rect rectBox;
            rectBox.x = left;
            rectBox.y = top;
            rectBox.width = width;
            rectBox.height = height;
            rectBoxes.push_back(rectBox);
            confidences.push_back(maxScore);
            class_ids.push_back(id);
        }
    }

    std::vector<int> indices;
    cv::dnn::NMSBoxes(rectBoxes, confidences, conf_threshold, nms_threshold, indices);

    for (size_t i = 0; i < indices.size(); i++) 
    {
        YoloBox yoloBox;
        int index = indices[i];
        yoloBox.id = class_ids[index];
        yoloBox.x = rectBoxes[index].x;
        yoloBox.y = rectBoxes[index].y;
        yoloBox.w = rectBoxes[index].width;
        yoloBox.h = rectBoxes[index].height;
        yoloBox.prob = confidences[index];
        yoloBoxes.push_back(yoloBox);
    }

    pthread_mutex_unlock(&yoloMutex);

    return yoloBoxes;
}

std::vector<YoloBox> Yolo::Detect(cv::Mat &img)
{
    return Detect(img, 0.25);
}

cv::Mat Yolo::Letterbox(cv::Mat& img, std::vector<float>& paddings, std::vector<int> new_shape)
{
    int img_h = img.rows; 
    int img_w = img.cols;

    // Compute scale ratio(new / old) and target resized shape
    float scale = std::min(new_shape[1] * 1.0 / img_h, new_shape[0] * 1.0 / img_w);
    int resize_h = int(round(img_h * scale));
    int resize_w = int(round(img_w * scale));
    paddings[0] = scale;

    // Compute padding
    int pad_h = new_shape[1] - resize_h;
    int pad_w = new_shape[0] - resize_w;

    // Resize and pad image while meeting stride-multiple constraints
    cv::Mat resized_img;
    cv::resize(img, resized_img, cv::Size(resize_w, resize_h));

    // divide padding into 2 sides
    float half_h = pad_h * 1.0 / 2;
    float half_w = pad_w * 1.0 / 2;
    paddings[1] = half_h;
    paddings[2] = half_w;

    // Compute padding boarder
    int top = int(round(half_h - 0.1));
    int bottom = int(round(half_h + 0.1));
    int left = int(round(half_w - 0.1));
    int right = int(round(half_w + 0.1));

    // Add border
    cv::copyMakeBorder(resized_img, resized_img, top, bottom, left, right, 0, cv::Scalar(114, 114, 114));

    return resized_img;
}