#include "Arith_BATask.h"
#include "ceres/ceres.h"
using namespace ceres;
// 定义残差结构体
struct HomographyResidual 
{
    HomographyResidual(const cv::KeyPoint& keypoint_i, const cv::KeyPoint& keypoint_j)
        : keypoint_i_(keypoint_i), keypoint_j_(keypoint_j) {}

    template <typename T>
    bool operator()(const T* const h_i, const T* const h_j, T* residual) const 
    {
        // 残差计算逻辑
        T H_i[9] = { h_i[0], h_i[1], h_i[2],
                    h_i[3], h_i[4], h_i[5],
                    h_i[6], h_i[7], T(1.0) };

        T H_j[9] = { h_j[0], h_j[1], h_j[2],
                    h_j[3], h_j[4], h_j[5],
                    h_j[6], h_j[7], T(1.0) };

        T p_i[3] = { T(keypoint_i_.pt.x), T(keypoint_i_.pt.y), T(1.0) };
        T p_j[3] = { T(keypoint_j_.pt.x), T(keypoint_j_.pt.y), T(1.0) };

        T P_i[3] = { T(0), T(0), T(0) };
        T P_j[3] = { T(0), T(0), T(0) };

        for (int row = 0; row < 3; row++) 
        {
            for (int col = 0; col < 3; col++) 
            {
                P_i[row] += H_i[row * 3 + col] * p_i[col];
                P_j[row] += H_j[row * 3 + col] * p_j[col];
            }
        }

        P_i[0] /= P_i[2];
        P_i[1] /= P_i[2];
        P_j[0] /= P_j[2];
        P_j[1] /= P_j[2];

        residual[0] = P_i[0] - P_j[0];
        residual[1] = P_i[1] - P_j[1];

        return true;
    }


private:
    const cv::KeyPoint keypoint_i_; // 第 i 帧图像中的特征点
    const cv::KeyPoint keypoint_j_; // 第 j 帧图像中的特征点
};



BA_Task::BA_Task(GeoStitcher * pGeoTrans)
{
    _GeoStitcher = pGeoTrans;
    _FeaMatcher = new FeatureMatcher(DetectorType::SIFT, MatcherType::FLANN);
    _imgVec.reserve(100);
}

BA_Task::~BA_Task()
{
}

SINT32 BA_Task::addFrame(GD_VIDEO_FRAME_S img, FrameInfo para)
{
    cv::Mat frame = Mat(img.u32Height, img.u32Width, CV_8UC1, (void*)img.u64VirAddr[0]);

    // 缓存图像帧
    if (img.enPixelFormat == GD_PIXEL_FORMAT_E::GD_PIXEL_FORMAT_GRAY_Y8)
    {
        _imgVec.emplace_back(frame.clone());
    }
    // 缓存外参
    _paraVec.push_back(para);

    // 缓存初始H
    Proj t_Proj = _GeoStitcher->AnlayseTform(para);
    auto H = _GeoStitcher->findHomography(t_Proj,_panPara);
    _origMatrix.push_back(H.clone());
    _currMatrix.push_back(H.clone());


    // 提取特征点
    std::vector<cv::KeyPoint> keypoints;
    cv::Mat descriptors;
    _FeaMatcher->extractFeatures(frame,keypoints,descriptors);
    _FeaPtVec.push_back(keypoints);
    _FeaDespVec.push_back(descriptors);

    // 缓存包围多边形
    _polygon.push_back(warpRectWithH(H,cv::Size(img.u32Width,img.u32Height)));

    return _imgVec.size();
}

void BA_Task::setPanPara(PanInfo info)
{
    _panPara = info;
}


void BA_Task::Test()
{
    
}

//#define SHOW_MATCH
void BA_Task::optimizeBA()
{

    remap(_origMatrix, "_origMatrix");

    // 计算匹配性矩阵
    CalMatchMat(0.3);

    google::InitGoogleLogging("ceres");

    // 将 cv::Mat 转换为 Ceres 需要的数组形式
    std::vector<double*> h_list(_currMatrix.size());


    for (int i = 0; i < _currMatrix.size(); i++)
    {
        h_list[i] = (double*)_currMatrix[i].data;
    }
    // 创建 Ceres 问题
    ceres::Problem problem;

    // 添加残差块
    int nParaCnt = 0;//参数组数

    for (int i = 0; i < _MatchMat.cols; i++)
    {
        for (int j = i + 1; j < _MatchMat.rows; j++)
        {
            // IOU满足条件才匹配特征点
            if (_IOUMat.at<float>(i, j) < 0.3)
            {
                continue;
            }
#ifdef SHOW_MATCH
            cv::Mat image(1000, 1000, CV_8UC3, cv::Scalar(0, 0, 0));
            cv::Mat imagetmp(1000, 1000, CV_8UC3, cv::Scalar(0, 0, 0));
            vector<vector<cv::Point2f>> tmpPoly;
            tmpPoly.push_back(_polygon[i]);
            tmpPoly.push_back(_polygon[j]);

            cv::warpPerspective(_imgVec[i], imagetmp, _origMatrix[i], imagetmp.size());

            // 生成遮罩（全白图像，表示有效区域）
            cv::Mat mask1 = cv::Mat::ones(_imgVec[i].size(), CV_8UC1) * 255;
            cv::Mat warped_mask1;
            cv::warpPerspective(mask1, warped_mask1, _origMatrix[i], image.size());

            imagetmp.copyTo(image, warped_mask1);

            cv::warpPerspective(_imgVec[j], imagetmp, _origMatrix[j], imagetmp.size());
            cv::Mat mask2 = cv::Mat::ones(_imgVec[i].size(), CV_8UC1) * 255;
            cv::Mat warped_mask2;
            cv::warpPerspective(mask2, warped_mask2, _origMatrix[j], image.size());

            imagetmp.copyTo(image, warped_mask2);

            drawPolygons(image, tmpPoly);

            // 显示绘制结果
            cv::imshow("Polygons", image);
            cv::waitKey(1);

#endif
            std::vector<cv::DMatch> matches;
            //_FeaMatcher->matchFeatures(_FeaDespVec[i],_FeaDespVec[j],matches);
            _FeaMatcher->matchFeatures_WithH(_FeaPtVec[i], _FeaDespVec[i], _FeaPtVec[j], _FeaDespVec[j], _origMatrix[i], _origMatrix[j], matches);

            // 图像特征匹配点对超过N对才认为有效
            if (matches.size() > 50)
            {
                _MatchMat.at<int>(i, j) = matches.size();
                _MatchMat.at<int>(j, i) = matches.size();
            }
            else
            {
                continue;
            }

#ifdef SHOW_MATCH
            // 可视化匹配结果
            cv::Mat img_matches;
            cv::drawMatches(
                _imgVec[i], _FeaPtVec[i], // 第一幅图像及其特征点
                _imgVec[j], _FeaPtVec[j], // 第二幅图像及其特征点
                matches,          // 匹配结果
                img_matches,      // 输出图像
                cv::Scalar::all(-1), // 匹配线颜色（默认随机颜色）
                cv::Scalar::all(-1), // 未匹配点颜色（默认不绘制）
                std::vector<char>(), // 掩码（可选，用于筛选匹配）
                cv::DrawMatchesFlags::NOT_DRAW_SINGLE_POINTS // 不绘制未匹配的点
            );

            cv::resize(img_matches, img_matches, cv::Size(1280, 512));

            // 显示匹配结果
            cv::imshow("Feature Matches", img_matches);
            cv::waitKey(0);
#endif

            // 添加匹配点对的残差块
            for (int m = 0; m < matches.size(); m++)
            {
                auto mc = matches[m];
                // 注意：这里不对，应该找匹配点！！ todo 节后完成
                cv::KeyPoint keypoint_i = _FeaPtVec[i][mc.queryIdx];
                cv::KeyPoint keypoint_j = _FeaPtVec[j][mc.trainIdx];

                ceres::LossFunction* loss_function = new ceres::HuberLoss(15);

                ceres::CostFunction* cost_function =
                    new ceres::AutoDiffCostFunction<HomographyResidual, 2, 8, 8>(
                        new HomographyResidual(keypoint_i, keypoint_j));
                problem.AddResidualBlock(cost_function, loss_function, h_list[i], h_list[j]);
            }

            nParaCnt += matches.size();
            
        }
    }



    // 配置求解器
    ceres::Solver::Options options;
    options.max_num_iterations = 2; // 增加最大迭代次数
    options.function_tolerance = 1e-8; // 设置更严格的函数值容忍度
    options.gradient_tolerance = 1e-10; // 设置更严格的梯度容忍度
    options.parameter_tolerance = 1e-10; // 设置更严格的参数容忍度
    options.minimizer_progress_to_stdout = true;
    //options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY; // 使用稀疏求解器
    options.num_threads = 16; // 使用多线程
    ceres::Solver::Summary summary;

    // 求解
    ceres::Solve(options, &problem, &summary);

    //// 将优化后的参数转换回 cv::Mat
    for (int i = 0; i < _currMatrix.size(); i++) 
    {
        std::cout << "------------" << std::endl;
        std::cout << _origMatrix[i] << std::endl;
        std::cout << _currMatrix[i] << std::endl;
        std::cout << "------------" << std::endl;
    }

    // 输出结果
    std::cout << summary.BriefReport() << std::endl;


    remap(_currMatrix, "_currMatrix");


    cv::waitKey(0);
}

void BA_Task::remap(vector<cv::Mat_<double>> H_Vec, std::string winname)
{
    cv::Mat image(1500, 1500, CV_8UC3, cv::Scalar(0, 0, 0));

    for (size_t i = 0; i < _imgVec.size(); i++)
    {
        cv::Mat imagetmp(1500, 1500, CV_8UC3, cv::Scalar(0, 0, 0));
        cv::warpPerspective(_imgVec[i], imagetmp, H_Vec[i], imagetmp.size());

        cv::Mat mask = cv::Mat::ones(_imgVec[i].size(), CV_8UC1) * 255;
        cv::Mat warped_mask;
        cv::warpPerspective(mask, warped_mask, H_Vec[i], image.size());
        imagetmp.copyTo(image, warped_mask);
    }

    imshow(winname, image);
    cv::waitKey(1);

}



//#define SHOW_MATCH
SINT32 BA_Task::CalMatchMat(float fiou_thre)
{
    _IOUMat = cv::Mat::zeros(_polygon.size(),_polygon.size(),CV_32FC1);

    _MatchMat = cv::Mat::zeros(_polygon.size(), _polygon.size(), CV_32SC1);


    // 先计算IOU矩阵
    for (size_t i = 0; i < _polygon.size(); i++)
    {
        vector<cv::Point2f> poly_i = _polygon[i];

        for (size_t j = i + 1; j < _polygon.size(); j++)
        {
            if (i == j)
            {
                _IOUMat.at<float>(i,j) = 1; 
                continue;
            }
            vector<cv::Point2f> poly_j = _polygon[j];

            float fiou = computeQuadrilateralIOU(poly_i,poly_j);
            _IOUMat.at<float>(i,j) = fiou;
            _IOUMat.at<float>(j,i) = fiou;//是对称矩阵
            
        }
    }


    return 0;
}




// 函数：绘制多边形
void drawPolygons(cv::Mat& image, const std::vector<std::vector<cv::Point2f>>& polygons) 
{
    // 定义颜色列表
    std::vector<cv::Scalar> colors = {
        cv::Scalar(255, 0, 0),   // 蓝色
        cv::Scalar(0, 255, 0),   // 绿色
        cv::Scalar(0, 0, 255),   // 红色
        cv::Scalar(255, 255, 0), // 青色
        cv::Scalar(255, 0, 255), // 品红
        cv::Scalar(0, 255, 255)  // 黄色
    };

    // 绘制每个多边形
    for (size_t i = 0; i < polygons.size(); ++i) {
        const auto& polygon = polygons[i];
        cv::Scalar color = colors[i % colors.size()]; // 循环使用颜色

        // 将多边形点转换为整数类型
        std::vector<cv::Point> intPolygon;
        for (const auto& pt : polygon) {
            intPolygon.push_back(cv::Point(static_cast<int>(pt.x), static_cast<int>(pt.y)));
        }

        // 绘制多边形
        cv::polylines(image, intPolygon, true, color, 2);
    }
}