#include "ChessboardDetector.h"
#include <opencv2/opencv.hpp>
#include <cmath>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

ChessboardDetector::ChessboardDetector()
    : m_adaptiveThresh(true)
    , m_normalizeImage(true)
    , m_filterQuads(false)
{
}

ChessboardDetector::~ChessboardDetector()
{
}

int ChessboardDetector::DetectChessboard(
    const unsigned char* imageData,
    int width,
    int height,
    int channels,
    int patternWidth,
    int patternHeight,
    ChessboardDetectResult& result)
{
    if (!imageData || width <= 0 || height <= 0) {
        return -1;
    }

    result.detected = false;
    result.corners.clear();
    result.patternWidth = patternWidth;
    result.patternHeight = patternHeight;
    result.hasPose = false;

    try {
        // 创建 OpenCV Mat
        cv::Mat image;
        if (channels == 1) {
            image = cv::Mat(height, width, CV_8UC1, (void*)imageData).clone();
        } else if (channels == 3) {
            image = cv::Mat(height, width, CV_8UC3, (void*)imageData).clone();
            cv::cvtColor(image, image, cv::COLOR_BGR2GRAY);
        } else {
            return -2;
        }

        // 检测标定板角点
        cv::Size patternSize(patternWidth, patternHeight);
        std::vector<cv::Point2f> corners;

        int flags = 0;
        if (m_adaptiveThresh) flags |= cv::CALIB_CB_ADAPTIVE_THRESH;
        if (m_normalizeImage) flags |= cv::CALIB_CB_NORMALIZE_IMAGE;
        if (m_filterQuads) flags |= cv::CALIB_CB_FILTER_QUADS;

        bool found = cv::findChessboardCorners(image, patternSize, corners, flags);

        if (found && corners.size() == patternWidth * patternHeight) {
            // 亚像素精细化
            cv::cornerSubPix(image, corners, cv::Size(11, 11), cv::Size(-1, -1),
                cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 30, 0.1));

            result.detected = true;
            result.corners.reserve(corners.size());
            for (const auto& pt : corners) {
                result.corners.push_back(Point2D(pt.x, pt.y));
            }
        }

        return 0;
    }
    catch (const cv::Exception& e) {
        return -3;
    }
}

int ChessboardDetector::DetectChessboardWithPose(
    const unsigned char* imageData,
    int width,
    int height,
    int channels,
    int patternWidth,
    int patternHeight,
    double squareSize,
    const CameraIntrinsics& intrinsics,
    ChessboardDetectResult& result)
{
    // 首先检测角点
    int ret = DetectChessboard(imageData, width, height, channels,
                               patternWidth, patternHeight, result);
    if (ret != 0 || !result.detected) {
        return ret;
    }

    try {
        // 构建3D物体点（标定板坐标系）
        std::vector<cv::Point3f> objectPoints;
        for (int i = 0; i < patternHeight; i++) {
            for (int j = 0; j < patternWidth; j++) {
                objectPoints.push_back(cv::Point3f(j * squareSize, i * squareSize, 0));
            }
        }

        // 构建2D图像点
        std::vector<cv::Point2f> imagePoints;
        for (const auto& pt : result.corners) {
            imagePoints.push_back(cv::Point2f(pt.x, pt.y));
        }

        // 相机内参矩阵
        cv::Mat cameraMatrix = (cv::Mat_<double>(3, 3) <<
            intrinsics.fx, 0, intrinsics.cx,
            0, intrinsics.fy, intrinsics.cy,
            0, 0, 1);

        // 畸变系数
        cv::Mat distCoeffs = (cv::Mat_<double>(5, 1) <<
            intrinsics.k1, intrinsics.k2, intrinsics.p1, intrinsics.p2, intrinsics.k3);

        // 求解PnP
        cv::Mat rvec, tvec;
        bool success = cv::solvePnP(objectPoints, imagePoints, cameraMatrix, distCoeffs,
                                     rvec, tvec, false, cv::SOLVEPNP_ITERATIVE);

        if (success) {
            result.hasPose = true;

            // 保存旋转向量和平移向量
            result.rvec[0] = rvec.at<double>(0);
            result.rvec[1] = rvec.at<double>(1);
            result.rvec[2] = rvec.at<double>(2);

            result.tvec[0] = tvec.at<double>(0);
            result.tvec[1] = tvec.at<double>(1);
            result.tvec[2] = tvec.at<double>(2);

            // 标定板中心位置（平移向量就是标定板原点在相机坐标系下的位置）
            result.center.x = tvec.at<double>(0);
            result.center.y = tvec.at<double>(1);
            result.center.z = tvec.at<double>(2);

            // 将旋转向量转换为旋转矩阵
            cv::Mat rotMat;
            cv::Rodrigues(rvec, rotMat);

            for (int i = 0; i < 3; i++) {
                for (int j = 0; j < 3; j++) {
                    result.rotationMatrix[i * 3 + j] = rotMat.at<double>(i, j);
                }
            }

            // 标定板法向量（标定板坐标系的Z轴在相机坐标系下的方向）
            // 旋转矩阵的第三列就是标定板Z轴方向
            result.normal.x = rotMat.at<double>(0, 2);
            result.normal.y = rotMat.at<double>(1, 2);
            result.normal.z = rotMat.at<double>(2, 2);

            // 计算欧拉角 (ZYX顺序)
            double sy = sqrt(rotMat.at<double>(0, 0) * rotMat.at<double>(0, 0) +
                           rotMat.at<double>(1, 0) * rotMat.at<double>(1, 0));

            bool singular = sy < 1e-6;

            double roll, pitch, yaw;
            if (!singular) {
                roll = atan2(rotMat.at<double>(2, 1), rotMat.at<double>(2, 2));
                pitch = atan2(-rotMat.at<double>(2, 0), sy);
                yaw = atan2(rotMat.at<double>(1, 0), rotMat.at<double>(0, 0));
            } else {
                roll = atan2(-rotMat.at<double>(1, 2), rotMat.at<double>(1, 1));
                pitch = atan2(-rotMat.at<double>(2, 0), sy);
                yaw = 0;
            }

            // 转换为度
            result.eulerAngles[0] = roll * 180.0 / M_PI;
            result.eulerAngles[1] = pitch * 180.0 / M_PI;
            result.eulerAngles[2] = yaw * 180.0 / M_PI;
        }

        return 0;
    }
    catch (const cv::Exception& e) {
        return -4;
    }
}

int ChessboardDetector::SetDetectionFlags(
    bool adaptiveThresh,
    bool normalizeImage,
    bool filterQuads)
{
    m_adaptiveThresh = adaptiveThresh;
    m_normalizeImage = normalizeImage;
    m_filterQuads = filterQuads;
    return 0;
}

// 工厂函数实现
IChessboardDetector* CreateChessboardDetectorInstance()
{
    return new ChessboardDetector();
}

void DestroyChessboardDetectorInstance(IChessboardDetector* instance)
{
    if (instance) {
        delete instance;
    }
}