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

using namespace std;

static tuple<cv::Mat, int, ShapeInfo> adjust_image_impl(const cv::Mat& img, const ShapeInfo& shapeinfo, int rotate_range) {
    int final_angle = 0;
    cv::Mat final_img;

    int img_h = img.rows;
    int img_w = img.cols;
    string shapetype = shapeinfo.type;
    cv::RotatedRect minAreaRect = shapeinfo.minAreaRect;
    cv::Point2f center = minAreaRect.center;
    cv::Size2f wh = minAreaRect.size;
    double wh_ratio = max(wh.width / wh.height, wh.height / wh.width);

    cv::Mat gray;
    cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);

    cv::Mat binary;
    cv::threshold(gray, binary, 200, 255, cv::THRESH_BINARY_INV | cv::THRESH_OTSU);

    if (shapetype == "ellipse" || shapetype == "circle") {
        binary = wipe_star(binary, minAreaRect);
    }
    binary = wipe_noise(binary);

    ShapeInfo new_shapeinfo = shapeinfo.copy();
    cv::Mat fck_gray = wipe_noise_gray(gray);

    if (shapetype == "ellipse" || shapetype == "circle") {
        fck_gray = my_dilate(fck_gray, cv::Size(16, 16), 1, WHITE_BG);
    }

    double min_cal = 1e10;
    cv::Mat first_M;
    int first_angle = 0;

    for (int angle = -rotate_range; angle < rotate_range; ++angle) {
        cv::Mat M = cv::getRotationMatrix2D(center, angle, 1.0);
        cv::Mat fck_gray1;
        cv::warpAffine(fck_gray, fck_gray1, M, cv::Size(img_w, img_h), cv::INTER_LINEAR, cv::BORDER_CONSTANT, cv::Scalar(255));

        int center_x = static_cast<int>(center.x);
        center_x = min(max(1, center_x), img_w - 1);

        int left_w = center_x;
        int right_start = (img_w % 2 == 0) ? center_x : (center_x + 1);
        int right_w = img_w - right_start;
        int min_w = min(left_w, right_w);

        cv::Mat left, right;
        if (min_w > 0) {
            left = fck_gray1(cv::Rect(0, 0, min_w, img_h)).clone();
            cv::Mat right_part = fck_gray1(cv::Rect(right_start, 0, min_w, img_h));
            cv::flip(right_part, right, 1);
        } else {
            left = cv::Mat::zeros(img_h, 1, fck_gray1.type());
            right = cv::Mat::zeros(img_h, 1, fck_gray1.type());
        }

        cv::Mat diff;
        cv::absdiff(left, right, diff);
        int cal = cv::countNonZero(diff);

        if (cal < min_cal) {
            min_cal = cal;
            first_M = M.clone();
            first_angle = angle;
        }
    }

    cv::Mat first_gray;
    cv::warpAffine(gray, first_gray, first_M, cv::Size(img_w, img_h), cv::INTER_LINEAR, cv::BORDER_CONSTANT, cv::Scalar(255));
    final_angle = first_angle;

    if (wh_ratio < 1.1 && abs(first_angle) > 45) {
        cv::Mat binary2;
        cv::threshold(first_gray, binary2, 200, 255, cv::THRESH_BINARY_INV | cv::THRESH_OTSU);

        int x_range[2] = {static_cast<int>(center.x - wh.width / 2 * 0.4),
                          static_cast<int>(center.x + wh.width / 2 * 0.4)};
        int y_range[2] = {static_cast<int>(center.y - wh.height / 2 * 0.85),
                          static_cast<int>(center.y + wh.height / 2 * 0.85)};

        int tx1 = max(0, x_range[0]);
        int ty1 = max(0, y_range[0]);
        int tx2 = min(static_cast<int>(first_gray.cols), x_range[1]);
        int ty2 = min(static_cast<int>(first_gray.rows), y_range[1]);
        
        if (tx1 < tx2 && ty1 < ty2) {
            int range_w = tx2 - tx1;
            int range_h = ty2 - ty1;
            if (range_w > 0 && range_h > 0) {
                cv::Mat hahaha = binary2(cv::Rect(tx1, ty1, range_w, range_h)).clone();
                
                int up_h = max(1, static_cast<int>(hahaha.rows / 4.0));
                int down_y = static_cast<int>(hahaha.rows / 4.0 * 3);
                down_y = min(max(down_y, 0), hahaha.rows - up_h);
                
                if (up_h > 0 && down_y >= 0 && down_y + up_h <= hahaha.rows) {
                    cv::Mat up = hahaha(cv::Rect(0, 0, hahaha.cols, up_h)).clone();
                    cv::Mat down = hahaha(cv::Rect(0, down_y, hahaha.cols, up_h)).clone();

                    int up_cnt = cv::countNonZero(up);
                    int down_cnt = cv::countNonZero(down);
                    double up_div_down = max(up_cnt, down_cnt) / max(1, min(up_cnt, down_cnt));

                    if (up_div_down > 2) {
                        if (up_cnt < down_cnt) {
                            final_angle = first_angle + 180;
                        }
                    }
                }
            }
        }
    }

    cv::Mat final_M = cv::getRotationMatrix2D(center, final_angle, 1);
    cv::warpAffine(img, final_img, final_M, cv::Size(img_w, img_h), cv::INTER_LINEAR, cv::BORDER_CONSTANT, cv::Scalar(255, 255, 255));

    new_shapeinfo.minAreaRect = cv::RotatedRect(new_shapeinfo.minAreaRect.center,
                                                new_shapeinfo.minAreaRect.size, 0);

    return make_tuple(final_img, final_angle, new_shapeinfo);
}

tuple<cv::Mat, int, ShapeInfo> adjust_img(const cv::Mat& img, const ShapeInfo& shapeinfo) {
    string shapetype = shapeinfo.type;

    if (shapetype == "ellipse" || shapetype == "circle") {
        return adjust_image_impl(img, shapeinfo, 90);
    } else if (shapetype == "square") {
        return adjust_image_impl(img, shapeinfo, 30);
    } else {
        return make_tuple(img.clone(), 0, shapeinfo.copy());
    }
}