#include "yolov8.h"

#include <cmath>
#include <cstdio>
#include <vector>

int main()
{
    constexpr int channels = 5;
    constexpr int positions = 6;

    rknn_tensor_attr output_attr{};
    output_attr.n_dims = 3;
    output_attr.dims[0] = 1;
    output_attr.dims[1] = channels;
    output_attr.dims[2] = positions;
    output_attr.type = RKNN_TENSOR_FLOAT32;

    rknn_app_context_t context{};
    context.io_num.n_output = 1;
    context.output_attrs = &output_attr;
    context.model_width = 640;
    context.model_height = 640;

    // NCP: [center_x, center_y, width, height, class_0_probability].
    // Put the lower-scored box first in the original tensor so the sort must
    // preserve the association between each original box and its confidence.
    std::vector<float> output(channels * positions, 0.0f);
    auto set = [&](int channel, int position, float value) {
        output[channel * positions + position] = value;
    };
    set(0, 0, 100.0f);
    set(1, 0, 100.0f);
    set(2, 0, 40.0f);
    set(3, 0, 40.0f);
    set(4, 0, 0.30f);

    set(0, 5, 500.0f);
    set(1, 5, 500.0f);
    set(2, 5, 40.0f);
    set(3, 5, 40.0f);
    set(4, 5, 0.90f);

    rknn_output rknn_output_value{};
    rknn_output_value.want_float = 1;
    rknn_output_value.buf = output.data();

    letterbox_t letterbox{};
    letterbox.scale = 1.0f;
    object_detect_result_list results{};
    const int status = post_process(
        &context, &rknn_output_value, &letterbox, 0.25f, 0.45f, &results);

    if (status != 0 || results.count != 2)
    {
        std::fprintf(stderr, "expected two detections, status=%d count=%d\n", status, results.count);
        return 1;
    }
    if (results.results[0].box.left < 400 ||
        std::fabs(results.results[0].prop - 0.90f) > 0.0001f)
    {
        std::fprintf(
            stderr,
            "highest score was detached from its box: left=%d score=%.6f\n",
            results.results[0].box.left,
            results.results[0].prop);
        return 2;
    }
    if (results.results[1].box.left > 200 ||
        std::fabs(results.results[1].prop - 0.30f) > 0.0001f)
    {
        std::fprintf(
            stderr,
            "lower score was detached from its box: left=%d score=%.6f\n",
            results.results[1].box.left,
            results.results[1].prop);
        return 3;
    }

    std::puts("DETECT_SCORE_MAPPING_TEST_OK");
    return 0;
}