#include "yolo_seg_geometry.h"
#include <opencv2/imgproc.hpp>

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

namespace {

bool close_to(double actual, double expected, double tolerance = 1e-6) {
    return std::abs(actual - expected) <= tolerance;
}

} // namespace

int main() {
    const unsigned char mask[] = {
        0, 0, 0, 0, 0, 0,
        0, 255, 255, 255, 255, 0,
        0, 255, 255, 255, 255, 0,
        0, 255, 255, 255, 255, 0,
        0, 0, 0, 0, 0, 0,
    };
    const auto boundary = yolo_seg_geometry::mask_boundary(mask, 6, 5, 6, 10, 20);
    assert(!boundary.empty());

    const auto rect = yolo_seg_geometry::minimum_area_rect(boundary);
    assert(rect.valid);
    assert(close_to(rect.center_x, 12.5));
    assert(close_to(rect.center_y, 22.0));
    assert(close_to(rect.width, 3.0));
    assert(close_to(rect.height, 2.0));

    const auto circle = yolo_seg_geometry::minimum_enclosing_circle(boundary);
    assert(circle.valid);
    assert(close_to(circle.center_x, 12.5));
    assert(close_to(circle.center_y, 22.0));
    assert(close_to(circle.radius, std::sqrt(13.0) * 0.5));
    const auto mean = yolo_seg_geometry::mean_diameter_circle(boundary);
    assert(mean.valid && !mean.fitted);
    assert(close_to(mean.center_x, rect.center_x));
    assert(close_to(mean.center_y, rect.center_y));
    assert(close_to(mean.radius, 1.25)); // Not half the short side, nor enclosing radius.

    cv::Mat disk = cv::Mat::zeros(121, 121, CV_8U);
    cv::circle(disk, {60, 60}, 30, cv::Scalar(255), -1);
    auto fit = yolo_seg_geometry::robust_fit_circle(disk.data, disk.cols, disk.rows, int(disk.step), 10, 20);
    assert(fit.valid && fit.fitted);
    assert(close_to(fit.center_x, 70, .1) && close_to(fit.center_y, 80, .1));
    assert(close_to(fit.radius, 30, 1));
    // Disconnected noise must not replace or enlarge the largest-contour fit.
    disk.at<unsigned char>(2, 2) = 255;
    const auto noisy = yolo_seg_geometry::robust_fit_circle(disk.data, disk.cols, disk.rows, int(disk.step), 10, 20);
    assert(noisy.fitted && close_to(noisy.radius, fit.radius));
    cv::line(disk, {87, 60}, {103, 60}, cv::Scalar(255), 2);
    const auto spiked = yolo_seg_geometry::robust_fit_circle(disk.data, disk.cols, disk.rows, int(disk.step), 0, 0);
    assert(spiked.valid && spiked.fitted && close_to(spiked.radius, 30, 1));
    cv::Mat elongated = cv::Mat::zeros(50, 100, CV_8U);
    cv::rectangle(elongated, {10, 10}, {80, 20}, cv::Scalar(255), -1);
    const auto fallback = yolo_seg_geometry::robust_fit_circle(elongated.data, elongated.cols, elongated.rows, int(elongated.step), 0, 0);
    assert(fallback.valid && !fallback.fitted);
    assert(close_to(fallback.radius, std::sqrt(70.*70.+10.*10.)*.5, .001));
    const unsigned char point[] = {255};
    const auto single = yolo_seg_geometry::robust_fit_circle(point, 1, 1, 1, 5, 8);
    assert(single.valid && !single.fitted && close_to(single.center_x, 5) && close_to(single.center_y, 8));
    assert(!yolo_seg_geometry::robust_fit_circle(nullptr, 1, 1, 1, 0, 0).valid);
    const unsigned char empty[] = {0};
    assert(!yolo_seg_geometry::robust_fit_circle(empty, 1, 1, 1, 0, 0).valid);
    assert(!yolo_seg_geometry::mean_diameter_circle({}).valid);

    const std::vector<yolo_seg_geometry::Point> rotated_square = {
        {0.0, -2.0}, {2.0, 0.0}, {0.0, 2.0}, {-2.0, 0.0}};
    const auto rotated = yolo_seg_geometry::minimum_area_rect(rotated_square);
    assert(rotated.valid);
    assert(close_to(rotated.center_x, 0.0));
    assert(close_to(rotated.center_y, 0.0));
    assert(close_to(rotated.width, std::sqrt(8.0)));
    assert(close_to(rotated.height, std::sqrt(8.0)));
    return 0;
}