#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);

}

bool isElliptical(float x, float y, float a2, float b2) {
    return (x * x / a2 + y * y / b2) < 1;
}

uint32_t getPixelIndex(int row, int col) {
    // LOG(INFO) << "col " << col << ", row " << row;
    col = MAX(0, MIN(1280, col)) % 640;
    row = MAX(0, MIN(960, row)) % 480;
    // LOG(INFO) << "col " << col << ", row " << row;
    return row * 640 + col;
}

// 识别结果数据包含识别结果、时间序号、具体识别信息
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) {
                // 内外椭圆参数
                float a1 = iter->rect.width * 0.5;
                float b1 = iter->rect.height * 0.5;
                float a2 = iter->rect.width * 0.57;
                float b2 = iter->rect.height * 0.57;

                //
                int center_x = iter->rect.x + iter->rect.width * 0.5;
                int center_y = iter->rect.y + iter->rect.height * 0.5;
                int min_x = MAX(0, center_x - a2);
                int min_y = MAX(0, center_y - b2);
                float a12 = a1 * a1;
                float b12 = b1 * b1;
                float a22 = a2 * a2;
                float b22 = b2 * b2;

                for (int x_value = min_x; x_value < center_x; x_value++) {
                    auto t_x_value = x_value - center_x;
                    for (int y_value = min_y; y_value < center_y; y_value++) {
                        auto t_y_value = y_value - center_y;

                        if (!isElliptical(t_x_value, t_y_value, a22, b22)) {
                            continue;
                        }
                        auto x_value_sym = 2 * center_x - x_value;
                        auto y_value_sym = 2 * center_y - y_value;
                        if (isElliptical(t_x_value, t_y_value, a12, b12)) {
                            seg_results.mutable_boxes(device_index)->add_wheels()->set_point(getPixelIndex(y_value, x_value));
                            seg_results.mutable_boxes(device_index)->add_wheels()->set_point(getPixelIndex(y_value, x_value_sym));
                            seg_results.mutable_boxes(device_index)->add_wheels()->set_point(getPixelIndex(y_value_sym, x_value));
                            seg_results.mutable_boxes(device_index)->add_wheels()->set_point(getPixelIndex(y_value_sym, x_value_sym));
                            continue;
                        }
                        seg_results.mutable_boxes(device_index)->add_tires()->set_point(getPixelIndex(y_value, x_value));
                        seg_results.mutable_boxes(device_index)->add_tires()->set_point(getPixelIndex(y_value, x_value_sym));
                        seg_results.mutable_boxes(device_index)->add_tires()->set_point(getPixelIndex(y_value_sym, x_value));
                        seg_results.mutable_boxes(device_index)->add_tires()->set_point(getPixelIndex(y_value_sym, x_value_sym));
                    }
                }
            }
        }
        // cv::imshow("merge_mat", merge_mat);
        // cv::waitKey(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;
}