// User settings. Supply bbox_arr from one model's inference output.
array<int> target_label_id_arr = {0,1};
int pixel_distance_to_combine = 100;

// Input BoundingBox uses probability for model confidence.
// If your runtime uses confidence instead, replace bbox.probability below.
class CustomBBox {
    int label_id;
    string label_name;
    float confidence;
    int x;
    int y;
    int width;
    int height;
    int center_x;
    int center_y;
    int label_color_r;
    int label_color_g;
    int label_color_b;
    float sum_confidence;
    int freq_confidence;
};

// String label_id key -> array<CustomBBox>.
dictionary merged_bbox_dict;

class BoxMergeClass {
    void reset_box_merge_class() {
        merged_bbox_dict.deleteAll();
    }

    void box_merge(array<BoundingBox@> bbox_arr, array<int> target_label_id_arr,
                   int pixel_distance_to_combine) {
        if (bbox_arr.length() == 0) {
            return;
        }
        for (uint i = 0; i < bbox_arr.length(); i++) {
            BoundingBox@ a_bbox = bbox_arr[i];
            if (target_label_id_arr.find(a_bbox.label_id) < 0) {
                continue;
            }
            int label_id = a_bbox.label_id;
            string label_id_key = label_id + "";
            CustomBBox temp_custom_bbox = create_custom_bbox(a_bbox);
            if (merged_bbox_dict.exists(label_id_key) == false) {
                array<CustomBBox> custom_bbox_arr;
                custom_bbox_arr.insertLast(temp_custom_bbox);
                merged_bbox_dict.set(label_id_key, custom_bbox_arr);
                continue;
            }
            bool found = false;
            array<CustomBBox> existing_custom_bbox_arr;
            found = merged_bbox_dict.get(label_id_key, existing_custom_bbox_arr);
            if (found == false) {
                continue;
            }
            bool found_family = false;
            for (uint j = 0; j < existing_custom_bbox_arr.length(); j++) {
                CustomBBox existing_custom_bbox = existing_custom_bbox_arr[j];
                if (existing_custom_bbox.label_id != label_id) {
                    continue;
                }
                float distance = calc_bbox_distance(existing_custom_bbox, temp_custom_bbox);
                if (distance <= pixel_distance_to_combine) {
                    found_family = true;
                    CustomBBox family_box = combine_bboxes(existing_custom_bbox, temp_custom_bbox);
                    existing_custom_bbox_arr[j] = family_box;
                    merged_bbox_dict.set(label_id_key, existing_custom_bbox_arr);
                    break;
                }
            }
            if (found_family == false) {
                existing_custom_bbox_arr.insertLast(temp_custom_bbox);
                merged_bbox_dict.set(label_id_key, existing_custom_bbox_arr);
            }
        }
    }

    CustomBBox create_custom_bbox(BoundingBox@ bbox) {
        CustomBBox output;
        output.label_id = bbox.label_id;
        output.label_name = bbox.label_name;
        output.confidence = bbox.probability;
        output.x = bbox.x;
        output.y = bbox.y;
        output.width = bbox.width;
        output.height = bbox.height;
        output.center_x = bbox.x + (bbox.width / 2);
        output.center_y = bbox.y + (bbox.height / 2);
        output.label_color_r = bbox.label_color_r;
        output.label_color_g = bbox.label_color_g;
        output.label_color_b = bbox.label_color_b;
        output.sum_confidence = bbox.probability;
        output.freq_confidence = 1;
        return output;
    }

    float calc_bbox_distance(const CustomBBox &in box1, const CustomBBox &in box2) {
        int box_1_min_x = box1.x;
        int box_1_max_x = box1.x + box1.width;
        int box_1_min_y = box1.y;
        int box_1_max_y = box1.y + box1.height;
        int box_2_min_x = box2.x;
        int box_2_max_x = box2.x + box2.width;
        int box_2_min_y = box2.y;
        int box_2_max_y = box2.y + box2.height;
        if (box_1_max_x >= box_2_min_x && box_2_max_x >= box_1_min_x &&
            box_1_max_y >= box_2_min_y && box_2_max_y >= box_1_min_y) {
            return 0.0f;
        }
        float x_gap = 0.0f;
        float y_gap = 0.0f;
        if (box_1_max_x < box_2_min_x) {
            x_gap = float(box_2_min_x - box_1_max_x);
        } else if (box_2_max_x < box_1_min_x) {
            x_gap = float(box_1_min_x - box_2_max_x);
        }
        if (box_1_max_y < box_2_min_y) {
            y_gap = float(box_2_min_y - box_1_max_y);
        } else if (box_2_max_y < box_1_min_y) {
            y_gap = float(box_1_min_y - box_2_max_y);
        }
        return sqrt(x_gap * x_gap + y_gap * y_gap);
    }

    CustomBBox combine_bboxes(CustomBBox existing_box1, CustomBBox new_box2) {
        CustomBBox combined;
        int left = min(existing_box1.x, new_box2.x);
        int top = min(existing_box1.y, new_box2.y);
        int right = max(existing_box1.x + existing_box1.width, new_box2.x + new_box2.width);
        int bottom = max(existing_box1.y + existing_box1.height, new_box2.y + new_box2.height);
        combined.x = left;
        combined.y = top;
        combined.width = right - left;
        combined.height = bottom - top;
        combined.center_x = combined.x + combined.width / 2;
        combined.center_y = combined.y + combined.height / 2;
        combined.label_id = existing_box1.label_id;
        combined.label_name = existing_box1.label_name;
        combined.label_color_r = existing_box1.label_color_r;
        combined.label_color_g = existing_box1.label_color_g;
        combined.label_color_b = existing_box1.label_color_b;
        combined.confidence = -1;
        combined.sum_confidence = existing_box1.sum_confidence + new_box2.confidence;
        combined.freq_confidence = existing_box1.freq_confidence + 1;
        return combined;
    }

    void calc_average_confidence() {
        array<string> label_ids_arr = merged_bbox_dict.getKeys();
        for (uint i = 0; i < label_ids_arr.length(); i++) {
            string label_id_key = label_ids_arr[i];
            array<CustomBBox> custom_bbox_arr;
            bool found = merged_bbox_dict.get(label_id_key, custom_bbox_arr);
            if (found == false) {
                continue;
            }
            for (uint j = 0; j < custom_bbox_arr.length(); j++) {
                CustomBBox a_custom_bbox = custom_bbox_arr[j];
                if (a_custom_bbox.confidence == -1) {
                    a_custom_bbox.confidence = float(a_custom_bbox.sum_confidence) /
                                              float(a_custom_bbox.freq_confidence);
                    custom_bbox_arr[j] = a_custom_bbox;
                }
            }
            merged_bbox_dict.set(label_id_key, custom_bbox_arr);
        }
    }

    int count_custom_bbox() {
        int total_count = 0;
        array<string> label_ids_arr = merged_bbox_dict.getKeys();
        for (uint i = 0; i < label_ids_arr.length(); i++) {
            string label_id_key = label_ids_arr[i];
            array<CustomBBox> bbox_arr;
            bool found = merged_bbox_dict.get(label_id_key, bbox_arr);
            if (found == false) {
                continue;
            }
            for (uint j = 0; j < bbox_arr.length(); j++) {
                CustomBBox a_custom_bbox = bbox_arr[j];
                if (target_label_id_arr.find(a_custom_bbox.label_id) < 0) {
                    continue;
                }
                total_count++;
            }
        }
        return total_count;
    }
};

BoxMergeClass box_merge_class;