#include "yolo_cpu_image_utils.h"
#include "yolo_letterbox.h"
#include "yolo_onnx_input_utils.h"

#include <cassert>
#include <cmath>
#include <vector>

static yolo_image_view_v1_t make_image(const unsigned char* data, int width, int height, int stride, yolo_pixel_format_t format) {
    yolo_image_view_v1_t image{};
    image.struct_size = sizeof(image);
    image.version = YOLO_IMAGE_VIEW_ABI_VERSION;
    image.data = data;
    image.width = width;
    image.height = height;
    image.row_stride_bytes = stride;
    image.pixel_format = format;
    return image;
}

int main() {
    const unsigned char rgb[] = {1, 2, 3, 4, 5, 6};
    auto rgb_image = make_image(rgb, 2, 1, 6, YOLO_PIXEL_FORMAT_RGB888);
    assert(yolo_cpu_image::valid(&rgb_image));
    rgb_image.row_stride_bytes = 5;
    assert(!yolo_cpu_image::valid(&rgb_image));

    const unsigned char nv12[] = {235, 235, 235, 235, 90, 240};
    const unsigned char nv21[] = {235, 235, 235, 235, 240, 90};
    auto nv12_image = make_image(nv12, 2, 2, 2, YOLO_PIXEL_FORMAT_NV12);
    auto nv21_image = make_image(nv21, 2, 2, 2, YOLO_PIXEL_FORMAT_NV21);
    assert(yolo_cpu_image::valid(&nv12_image));
    assert(yolo_cpu_image::valid(&nv21_image));
    float r12 = 0, g12 = 0, b12 = 0, r21 = 0, g21 = 0, b21 = 0;
    yolo_cpu_image::pixel_rgb(nv12_image, 0, 0, r12, g12, b12);
    yolo_cpu_image::pixel_rgb(nv21_image, 0, 0, r21, g21, b21);
    assert(std::abs(r12 - r21) < 0.01f && std::abs(g12 - g21) < 0.01f && std::abs(b12 - b21) < 0.01f);
    assert(r12 > g12 && g12 > b12);

    assert(yolo_onnx_input::resolve_nchw({1, 3, 640, 640}, "nhwc"));
    assert(!yolo_onnx_input::resolve_nchw({1, 640, 640, 3}, "nchw"));
    assert(yolo_onnx_input::resolve_nchw({1, -1, -1, -1}, "nchw"));
    assert(!yolo_onnx_input::resolve_nchw({1, -1, -1, -1}, "nhwc"));
    yolo_onnx_input::validate_three_channels({1, 3, -1, -1}, true);
    yolo_onnx_input::validate_three_channels({1, -1, -1, 3}, false);

    const auto portrait = yolo_letterbox::make_plan(1348, 2120, 1280, 1280);
    assert(std::abs(portrait.scale - 1280.0f / 2120.0f) < 0.000001f);
    assert(portrait.resized_width == 814 && portrait.resized_height == 1280);
    assert(portrait.pad_x == 233 && portrait.pad_y == 0);
    assert(std::abs(yolo_letterbox::source_coordinate(0, 1348, 814) - 0.328010f) < 0.0001f);
    assert(yolo_letterbox::bilinear_u8(10, 20, 30, 40, 0.5f, 0.5f) == 25);
    return 0;
}