/**
 * @file correction.cpp
 * @brief 透视校正 — 从任意角度拍摄的数独照片中提取规整的 360×360 盘面
 *
 * 算法流程：
 *   1. 灰度化 + 高斯模糊 + 自适应二值化
 *   2. 轮廓检测 → 找最大正方形轮廓（数独盘面外框）
 *   3. 若轮廓不够方 → 膨胀连接碎片网格线 → 再次检测
 *   4. approxPolyDP 拟合四边形 → 对角点排序
 *   5. 透视变换 + resize → 360×360 规整盘面
 */

#include "correction.h"

#include <opencv2/imgproc.hpp>
#include <algorithm>
#include <cmath>
#include <stdexcept>
#include <vector>

namespace {

constexpr int SIZE_PUZZLE = 360;  // 输出尺寸 = 9 × 40px

// ── 正方形检查：宽高比 0.75~1.33，面积≥图像 15% ──
bool is_square_like(const std::vector<cv::Point>& contour, int img_area) {
    cv::Rect r = cv::boundingRect(contour);
    double area = cv::contourArea(contour);
    if (img_area <= 0 || area < img_area * 0.15) return false;
    if (r.width <= 0 || r.height <= 0) return false;
    double aspect = static_cast<double>(r.width) / r.height;
    return (aspect >= 0.75 && aspect <= 1.33);
}

// ── 轮廓选择：正方形优先，否则面积最大 ──
int select_best(const std::vector<std::vector<cv::Point>>& contours, int img_area) {
    int best_idx = 0, best_sq_idx = -1;
    double best_area = 0.0, best_sq_area = 0.0;
    for (size_t i = 0; i < contours.size(); ++i) {
        double area = cv::contourArea(contours[i]);
        if (area > best_area) { best_area = area; best_idx = static_cast<int>(i); }
        if (is_square_like(contours[i], img_area) && area > best_sq_area) {
            best_sq_area = area; best_sq_idx = static_cast<int>(i);
        }
    }
    return (best_sq_idx >= 0) ? best_sq_idx : best_idx;
}

// ── 四边形拟合：失败时回退到外接矩形 ──
std::vector<cv::Point> get_four_corners(const std::vector<cv::Point>& contour) {
    double arc_len = cv::arcLength(contour, true);
    std::vector<cv::Point> poly;

    cv::approxPolyDP(contour, poly, 0.1 * arc_len, true);
    if (poly.size() == 4) return poly;

    cv::approxPolyDP(contour, poly, 0.02 * arc_len, true);  // 更宽松的 epsilon
    if (poly.size() == 4) return poly;

    // 回退：直接用外接矩形的四个角
    cv::Rect r = cv::boundingRect(contour);
    return {{r.x, r.y}, {r.x + r.width, r.y},
            {r.x, r.y + r.height}, {r.x + r.width, r.y + r.height}};
}

// ── 角点排序：均值二分法 → 左上、右上、左下、右下 ──
std::vector<cv::Point2f> sort_corners(const std::vector<cv::Point>& corners) {
    assert(corners.size() == 4);
    float sum_x = 0, sum_y = 0;
    for (const auto& p : corners) { sum_x += p.x; sum_y += p.y; }
    float mx = sum_x / 4, my = sum_y / 4;

    std::vector<cv::Point2f> result(4);
    for (const auto& p : corners) {
        if (p.x < mx)
            result[(p.y < my) ? 0 : 2] = cv::Point2f(p.x, p.y);  // 左上:左下
        else
            result[(p.y < my) ? 1 : 3] = cv::Point2f(p.x, p.y);  // 右上:右下
    }
    return result;
}

}  // namespace

// ══════════════════════════════════════════════════════════════
//  correct_perspective — 对外接口
//  输入：任意尺寸的灰度图
//  输出：360×360 的规整盘面
// ══════════════════════════════════════════════════════════════
cv::Mat vision::correct_perspective(const cv::Mat& img_original) {
    // 1. 灰度化
    cv::Mat img_gray;
    if (img_original.channels() >= 3)
        cv::cvtColor(img_original, img_gray, cv::COLOR_BGR2GRAY);
    else
        img_gray = img_original;

    // 2. 高斯模糊 + 自适应二值化
    cv::Mat img_blur, img_thresh;
    cv::GaussianBlur(img_gray, img_blur, cv::Size(3, 3), 0);
    cv::adaptiveThreshold(img_blur, img_thresh, 255,
                          cv::ADAPTIVE_THRESH_GAUSSIAN_C,
                          cv::THRESH_BINARY_INV, 11, 2);
    int img_area = img_gray.rows * img_gray.cols;

    // 3. 第一遍轮廓检测（原始图像，角点更精确）
    std::vector<std::vector<cv::Point>> contours_orig;
    cv::findContours(img_thresh, contours_orig,
                     cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    if (contours_orig.empty())
        throw std::runtime_error("图像中未检测到任何轮廓，无法定位数独盘面");

    // 4. 若最佳轮廓不够正方形 → 膨胀网格线 → 第二遍检测
    int best_idx = select_best(contours_orig, img_area);
    std::vector<cv::Point> clean_contour;

    if (!is_square_like(contours_orig[best_idx], img_area)) {
        cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));
        cv::Mat dilated;
        cv::dilate(img_thresh, dilated, kernel);
        std::vector<std::vector<cv::Point>> contours_dilated;
        cv::findContours(dilated, contours_dilated,
                         cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
        if (contours_dilated.empty())
            throw std::runtime_error("图像中未检测到任何轮廓，无法定位数独盘面");
        clean_contour = contours_dilated[select_best(contours_dilated, img_area)];
    } else {
        clean_contour = contours_orig[best_idx];
    }

    // 5. 四边形拟合 → 角点排序
    auto sorted = sort_corners(get_four_corners(clean_contour));

    // 6. 透视变换到正方形
    float diag = std::min(sorted[3].x - sorted[0].x, sorted[3].y - sorted[0].y);
    int pic_size = static_cast<int>(diag);
    if (pic_size <= 0) pic_size = SIZE_PUZZLE;

    cv::Point2f src[4] = {sorted[0], sorted[1], sorted[2], sorted[3]};
    cv::Point2f dst[4] = {{0,0}, {(float)pic_size,0},
                          {0,(float)pic_size}, {(float)pic_size,(float)pic_size}};

    cv::Mat warp_mat = cv::getPerspectiveTransform(src, dst);
    cv::Mat warped;
    cv::warpPerspective(img_gray, warped, warp_mat, cv::Size(pic_size, pic_size));

    // 7. 最终模糊 + resize → 360×360
    cv::GaussianBlur(warped, warped, cv::Size(3, 3), 0);
    cv::resize(warped, warped, cv::Size(SIZE_PUZZLE, SIZE_PUZZLE));

    return warped;
}