#include "yolo_seg_mqtt_async.h"

void YoloSegmentMqttAsyncClient::init(const std::string &file) {
    m_message_arrived_callback_ = CloudDataArrived;

#ifdef ENABLE_TENSORRT_DETECT
    detector = new TensorrtWheelDetector(ETC_PATH PROJECT_NAME "/model.engine", ETC_PATH PROJECT_NAME "/class.txt");
#else
    detector = new TensorrtWheelDetector(ETC_PATH PROJECT_NAME "/model.onnx", ETC_PATH PROJECT_NAME "/class.txt");
#endif

    MqttAsyncClient::init_from_proto(file);

}

// 识别结果数据包含识别结果、时间序号、具体识别信息
int YoloSegmentMqttAsyncClient::CloudDataArrived(void *client, char *topicName, int topicLen,
                                                 MQTTAsync_message *message) {
    auto *tof_client = (YoloSegmentMqttAsyncClient *) client;
    // TODO：测试
    std::string topic = topicName;
    if (topic == "tof/ir") {
        // 数据反序列化
        int lable;
        cv::Mat merge_mat;
        JetStream::LabelImage testimage2;
        JetStream::LabelYolo seg_results;
        testimage2.ParseFromArray(message->payload, message->payloadlen);
        NetMessageTrans::Proto2lableMat(testimage2, lable, merge_mat);
        seg_results.set_label(lable);
        seg_results.add_boxes();
        seg_results.add_boxes();
        seg_results.add_boxes();
        seg_results.add_boxes();

        // 模型识别
        std::vector<Object> objs;
        cv::cvtColor(merge_mat, merge_mat, cv::COLOR_GRAY2RGB);
        tof_client->detector->detect(merge_mat, objs, merge_mat);
//        cv::imshow("merge_mat", merge_mat);
//        cv::waitKey(1);

        // 分离识别结果
        auto t = std::chrono::steady_clock::now();
        for (auto iter = objs.begin(); iter != objs.end(); iter++) {
            auto seg_points = tof_client->detector->getPointsFromObj(*iter);
            int device_index = (int(iter->rect.x / 640) * 0x01) | (int(iter->rect.y / 480) << 1);
            // 校验识别矩形框是否越界
            int device_index_check = (int((iter->rect.x + iter->rect.width) / 640) * 0x01) |
                                     (int((iter->rect.y + iter->rect.height) / 480) << 1);
            if (device_index != device_index_check) {
                // TODO:存图
                objs.erase(iter);
                iter--;
                continue;
            }
            int x_alpha = (device_index & 0x1);
            int y_alpha = ((device_index & 0x2) >> 1);
            seg_results.mutable_boxes(device_index)->set_x(iter->rect.x - x_alpha * merge_mat.cols / 2);
            seg_results.mutable_boxes(device_index)->set_y(iter->rect.y - y_alpha * merge_mat.rows / 2);
            seg_results.mutable_boxes(device_index)->set_width(iter->rect.width);
            seg_results.mutable_boxes(device_index)->set_height(iter->rect.height);
            seg_results.mutable_boxes(device_index)->set_confidence(iter->prob);
            if (iter->prob > 0.6) {
                // LOG(INFO) << device_index << ": " << iter->prob;
                int box_contour[480][2] = {0};
                for (auto &pt: seg_points) {    // use 7.5~9ms
                    pt.x -= x_alpha * merge_mat.cols / 2;
                    pt.y -= y_alpha * merge_mat.rows / 2;
                    if (box_contour[pt.y][0] == 0 && box_contour[pt.y][1] == 0) {
                        box_contour[pt.y][0] = pt.x;
                        box_contour[pt.y][1] = pt.x;
                    } else {
                        box_contour[pt.y][0] = MIN(pt.x, box_contour[pt.y][0]);
                        box_contour[pt.y][1] = MAX(pt.x, box_contour[pt.y][1]);
                    }
                }
                for (auto &ends: box_contour) {
                    auto line = seg_results.mutable_boxes(device_index)->add_lines();
                    line->set_begin(ends[0]);
                    line->set_end(ends[1]);
                }
            }
        }

        char data[seg_results.ByteSizeLong()];
        seg_results.SerializeToArray((void *) data, seg_results.ByteSizeLong());
        tof_client->SendMessage("tof/seg", data, seg_results.ByteSizeLong());
    }
    return 1;
}