#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include <opencv2/calib3d.hpp>

// ============================================================
// ==================== 需要修改的配置 =========================
// ============================================================

// ---- 输入 ----
#define INPUT_IMAGE       "G_V_p1_f94_1560157629850.jpeg"          // 待去畸变的原图路径
#define MAP_OUTPUT_FILE   "gdc_map.txt"        // 输出的 GDC map 文件
#define IMG_OUTPUT_FILE   "undistorted.png"    // 输出的去畸变图像
#define JSON_OUTPUT_FILE  "new_camera_matrix.json"  // 输出的新内参

// ---- 去畸变参数 ----
#define ALPHA             (0.0)   // 0=裁剪所有黑边, 1=保留所有像素, 0.5=平衡

// ============================================================
// ==================== 以下为功能实现 =========================
// ============================================================

int main()
{
    // ---------- 1. 相机内参 (来自 camera_port3.json) ----------
    cv::Mat cameraMatrix = (cv::Mat_<double>(3, 3) <<
        1074.073092471065,  0.0,               1028.482207839880,
        0.0,                1074.073092471065, 638.786153564850,
        0.0,                0.0,               1.0);

    // ---------- 2. 畸变系数 (8 个: k1,k2,p1,p2,k3,k4,k5,k6) ----------
    cv::Mat distCoeffs = (cv::Mat_<double>(8, 1) <<
      77.603041676467,
      63.081253715919,
      0.000115897214,
      0.000308477950,
      4.572721024064,
      77.860103887550,
      86.890135640424,
      15.880822580737); 

    // ---------- 3. 读取原图，获取尺寸 ----------
    cv::Mat img = cv::imread(INPUT_IMAGE);
    if (img.empty()) {
        std::cout << "Failed to read image: " << INPUT_IMAGE << std::endl;
        return -1;
    }
    int W = img.cols;
    int H = img.rows;
    cv::Size size(W, H);
    std::cout << "Image size: " << W << " x " << H << std::endl;

    // ---------- 4. 计算新的相机内参与去畸变映射 ----------
    cv::Rect roi;
    cv::Mat newCameraMatrix = cv::getOptimalNewCameraMatrix(
        cameraMatrix, distCoeffs, size, ALPHA, size, &roi, true);

    cv::Mat map1, map2;
    cv::initUndistortRectifyMap(
        cameraMatrix, distCoeffs, cv::Mat(),
        newCameraMatrix, size, CV_32FC1, map1, map2);

    std::cout << "new_camera_matrix:\n" << newCameraMatrix << std::endl;
    std::cout << "roi: x=" << roi.x << " y=" << roi.y
              << " w=" << roi.width << " h=" << roi.height << std::endl;

    // ---------- 5. 生成去畸变图像 ----------
    // cv::Mat undistorted;
    // cv::remap(img, undistorted, map1, map2,
    //           cv::INTER_LINEAR, cv::BORDER_CONSTANT);
    // cv::imwrite(IMG_OUTPUT_FILE, undistorted);
    // std::cout << "Undistorted image saved: " << IMG_OUTPUT_FILE << std::endl;

    // ---------- 6. 保存新的内参 JSON ----------
    std::ofstream jf(JSON_OUTPUT_FILE);
    if (jf.is_open()) {
        jf << "{\n";
        jf << "  \"source_config\": \"camera_port1.json\",\n";
        jf << "  \"camera_matrix\": {\n";
        jf << "    \"fx\": " << newCameraMatrix.at<double>(0, 0) << ",\n";
        jf << "    \"fy\": " << newCameraMatrix.at<double>(1, 1) << ",\n";
        jf << "    \"cx\": " << newCameraMatrix.at<double>(0, 2) << ",\n";
        jf << "    \"cy\": " << newCameraMatrix.at<double>(1, 2) << "\n";
        jf << "  },\n";
        jf << "  \"roi\": {\n";
        jf << "    \"x\": " << roi.x << ", \"y\": " << roi.y << ",\n";
        jf << "    \"w\": " << roi.width << ", \"h\": " << roi.height << "\n";
        jf << "  }\n";
        jf << "}\n";
        jf.close();
        std::cout << "New camera matrix saved: " << JSON_OUTPUT_FILE << std::endl;
    } else {
        std::cout << "Failed to open " << JSON_OUTPUT_FILE << std::endl;
    }

    // ---------- 7. 写出 gdc_map.txt ----------
    // GDC 工具的 map 文件是按 1280x720 头部格式书写的，这里保持原格式
    // 如果你的原图不是 1280x720，请把 CAM_WIDTH/CAM_HEIGHT 换成实际尺寸
    const int CAM_WIDTH  = W;
    const int CAM_HEIGHT = H;

    std::ofstream ofs(MAP_OUTPUT_FILE);
    if (!ofs.is_open()) {
        std::cout << "Failed to open " << MAP_OUTPUT_FILE << std::endl;
        return -1;
    }

    // 头部信息 (与原 GDC 工具格式一致)
    ofs << 1 << std::endl;
    ofs << 50 << " " << 50 << std::endl;
    ofs << CAM_HEIGHT << " " << CAM_WIDTH << std::endl;
    ofs << (CAM_HEIGHT / 2 - 1) << " " << (CAM_WIDTH / 2 - 1) << std::endl;

    // 逐像素写出映射: 先 y (map2) 再 x (map1)，用 ':' 分隔，像素间用空格
    for (int r = 0; r < CAM_HEIGHT; r++) {
        for (int c = 0; c < CAM_WIDTH; c++) {
            float x = map1.at<float>(r, c);
            float y = map2.at<float>(r, c);
            ofs << y << ":" << x << " ";
        }
        ofs << std::endl;
    }

    ofs.flush();
    ofs.close();
    std::cout << MAP_OUTPUT_FILE << " generated." << std::endl;

    return 0;
}