#include "image_morphology.h"
#include <opencv2/imgproc.hpp>
#include <cmath>
#include <algorithm>
#include <cassert>

using namespace std;

const string WHITE_BG = "white";
const string BLACK_BG = "black";

cv::Mat my_dilate(const cv::Mat& img, const cv::Size& ksize, int iter, const string& bg_type) {
    assert(bg_type == WHITE_BG || bg_type == BLACK_BG);
    
    cv::Mat kernel(ksize.height, ksize.width, CV_8UC1, cv::Scalar(255));
    cv::Mat result;
    
    if (bg_type == BLACK_BG) {
        cv::dilate(img, result, kernel, cv::Point(-1, -1), iter);
    } else {
        cv::erode(img, result, kernel, cv::Point(-1, -1), iter);
    }
    
    return result;
}

cv::Mat wipe_noise(const cv::Mat& img) {
    cv::Mat result = img.clone();
    int num_objects;
    cv::Mat labels, stats, centroids;
    num_objects = cv::connectedComponentsWithStats(img, labels, stats, centroids, 8);
    
    int img_min = min(img.rows, img.cols);
    double scale = pow(max(1.0, img_min / 300.0), 2);
    int threshold_area = static_cast<int>(10 * scale);
    
    for (int index = 0; index < num_objects; ++index) {
        int area = stats.at<int>(index, 4);
        if (area < threshold_area) {
            result.setTo(0, labels == index);
        }
    }
    return result;
}

cv::Mat wipe_noise_gray(const cv::Mat& img) {
    cv::Mat result = img.clone();
    cv::Mat blur;
    cv::GaussianBlur(img, blur, cv::Size(3, 3), 0);
    
    cv::Mat binary;
    cv::threshold(blur, binary, 200, 255, cv::THRESH_BINARY_INV | cv::THRESH_OTSU);
    
    int num_objects;
    cv::Mat labels, stats, centroids;
    num_objects = cv::connectedComponentsWithStats(binary, labels, stats, centroids, 8);
    
    int img_min = min(img.rows, img.cols);
    double scale = pow(max(1.0, img_min / 300.0), 2);
    int threshold_area = static_cast<int>(10 * scale);
    
    for (int index = 0; index < num_objects; ++index) {
        int area = stats.at<int>(index, 4);
        if (area < threshold_area) {
            result.setTo(255, labels == index);
        }
    }
    
    return result;
}

cv::Mat wipe_star(const cv::Mat& img, const cv::RotatedRect& minAreaRect) {
    cv::Mat result = img.clone();
    cv::Point2f center = minAreaRect.center;
    cv::Size2f wh = minAreaRect.size;
    double r = max(wh.width, wh.height) / min(wh.width, wh.height);
    int base_wh = static_cast<int>(wh.height / 2 + wh.height / 20);
    double _W = base_wh * r;
    double _H = base_wh;

    int num_objects;
    cv::Mat labels, stats, centroids;
    num_objects = cv::connectedComponentsWithStats(img, labels, stats, centroids, 8);

    for (int index = 1; index < num_objects; ++index) {
        int x = stats.at<int>(index, 0);
        int y = stats.at<int>(index, 1);
        int w = stats.at<int>(index, 2);
        int h = stats.at<int>(index, 3);
        cv::Point2f _center(x + w / 2.0, y + h / 2.0);
        double center_dis = sqrt(pow(_center.x - center.x, 2) + pow(_center.y - center.y, 2));
        double hehe = min(w, h);

        if (hehe > max(_W, _H) * 0.25 && hehe < min(_W, _H) && center_dis < max(_W, _H) * 0.2) {
            result.setTo(0, labels == index);
        }
    }

    if (r < 1.2) {
        int y1 = static_cast<int>(center.y - wh.height * 0.15);
        int y2 = static_cast<int>(center.y + wh.height * 0.15);
        int x1 = static_cast<int>(center.x - wh.width * 0.15);
        int x2 = static_cast<int>(center.x + wh.width * 0.15);
        
        int tx1 = max(0, x1);
        int ty1 = max(0, y1);
        int tx2 = min(static_cast<int>(result.cols), x2);
        int ty2 = min(static_cast<int>(result.rows), y2);
        
        if (tx1 < tx2 && ty1 < ty2) {
            result(cv::Rect(tx1, ty1, tx2 - tx1, ty2 - ty1)).setTo(0);
        }
    }

    return result;
}