GrabBag/Tools/CalibView/Src/BatchVerifyDialog.cpp

#include "BatchVerifyDialog.h"
#include "IChessboardDetector.h"
#include "IHandEyeCalib.h"

#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QGridLayout>
#include <QGroupBox>
#include <QFileDialog>
#include <QMessageBox>
#include <QHeaderView>
#include <QDir>
#include <QFileInfo>
#include <QSettings>
#include <QDateTime>
#include <QApplication>
#include <QImage>
#include <QJsonDocument>
#include <QJsonObject>
#include <QJsonArray>
#include <cmath>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

BatchVerifyDialog::BatchVerifyDialog(IChessboardDetector* detector,
                                     IHandEyeCalib* calib,
                                     QWidget* parent)
    : QDialog(parent)
    , m_detector(detector)
    , m_calib(calib)
    , m_lblDirectory(nullptr)
    , m_btnSelectDir(nullptr)
    , m_btnStartVerify(nullptr)
    , m_btnStopVerify(nullptr)
    , m_btnExport(nullptr)
    , m_progressBar(nullptr)
    , m_tableResults(nullptr)
    , m_logEdit(nullptr)
    , m_sbPatternWidth(nullptr)
    , m_sbPatternHeight(nullptr)
    , m_sbSquareSize(nullptr)
    , m_sbFx(nullptr)
    , m_sbFy(nullptr)
    , m_sbCx(nullptr)
    , m_sbCy(nullptr)
    , m_isVerifying(false)
{
    setupUI();
    setWindowTitle("批量验证工具");
    resize(1000, 700);
}

BatchVerifyDialog::~BatchVerifyDialog()
{
}

void BatchVerifyDialog::setupUI()
{
    QVBoxLayout* mainLayout = new QVBoxLayout(this);

    // 目录选择区域
    QGroupBox* dirGroup = new QGroupBox("数据目录", this);
    QHBoxLayout* dirLayout = new QHBoxLayout(dirGroup);
    m_lblDirectory = new QLabel("未选择目录", this);
    m_btnSelectDir = new QPushButton("选择目录...", this);
    connect(m_btnSelectDir, &QPushButton::clicked, this, &BatchVerifyDialog::onSelectDirectory);
    dirLayout->addWidget(m_lblDirectory, 1);
    dirLayout->addWidget(m_btnSelectDir);
    mainLayout->addWidget(dirGroup);

    // 参数设置区域
    QGroupBox* paramGroup = new QGroupBox("检测参数", this);
    QGridLayout* paramLayout = new QGridLayout(paramGroup);

    // 标定板参数
    paramLayout->addWidget(new QLabel("标定板宽度:", this), 0, 0);
    m_sbPatternWidth = new QSpinBox(this);
    m_sbPatternWidth->setRange(3, 20);
    m_sbPatternWidth->setValue(11);
    paramLayout->addWidget(m_sbPatternWidth, 0, 1);

    paramLayout->addWidget(new QLabel("标定板高度:", this), 0, 2);
    m_sbPatternHeight = new QSpinBox(this);
    m_sbPatternHeight->setRange(3, 20);
    m_sbPatternHeight->setValue(8);
    paramLayout->addWidget(m_sbPatternHeight, 0, 3);

    paramLayout->addWidget(new QLabel("方格尺寸(mm):", this), 0, 4);
    m_sbSquareSize = new QDoubleSpinBox(this);
    m_sbSquareSize->setRange(1, 100);
    m_sbSquareSize->setDecimals(2);
    m_sbSquareSize->setValue(15.0);
    paramLayout->addWidget(m_sbSquareSize, 0, 5);

    // 相机内参
    paramLayout->addWidget(new QLabel("fx:", this), 1, 0);
    m_sbFx = new QDoubleSpinBox(this);
    m_sbFx->setRange(0, 10000);
    m_sbFx->setDecimals(3);
    m_sbFx->setValue(1000.0);
    paramLayout->addWidget(m_sbFx, 1, 1);

    paramLayout->addWidget(new QLabel("fy:", this), 1, 2);
    m_sbFy = new QDoubleSpinBox(this);
    m_sbFy->setRange(0, 10000);
    m_sbFy->setDecimals(3);
    m_sbFy->setValue(1000.0);
    paramLayout->addWidget(m_sbFy, 1, 3);

    paramLayout->addWidget(new QLabel("cx:", this), 1, 4);
    m_sbCx = new QDoubleSpinBox(this);
    m_sbCx->setRange(0, 5000);
    m_sbCx->setDecimals(3);
    m_sbCx->setValue(640.0);
    paramLayout->addWidget(m_sbCx, 1, 5);

    paramLayout->addWidget(new QLabel("cy:", this), 2, 0);
    m_sbCy = new QDoubleSpinBox(this);
    m_sbCy->setRange(0, 5000);
    m_sbCy->setDecimals(3);
    m_sbCy->setValue(512.0);
    paramLayout->addWidget(m_sbCy, 2, 1);

    mainLayout->addWidget(paramGroup);

    // 控制按钮
    QHBoxLayout* btnLayout = new QHBoxLayout();
    m_btnStartVerify = new QPushButton("开始验证", this);
    m_btnStopVerify = new QPushButton("停止", this);
    m_btnExport = new QPushButton("导出结果", this);
    m_btnStartVerify->setEnabled(false);
    m_btnStopVerify->setEnabled(false);
    m_btnExport->setEnabled(false);
    connect(m_btnStartVerify, &QPushButton::clicked, this, &BatchVerifyDialog::onStartVerify);
    connect(m_btnStopVerify, &QPushButton::clicked, this, &BatchVerifyDialog::onStopVerify);
    connect(m_btnExport, &QPushButton::clicked, this, &BatchVerifyDialog::onExportResults);
    btnLayout->addWidget(m_btnStartVerify);
    btnLayout->addWidget(m_btnStopVerify);
    btnLayout->addWidget(m_btnExport);
    btnLayout->addStretch();
    mainLayout->addLayout(btnLayout);

    // 进度条
    m_progressBar = new QProgressBar(this);
    m_progressBar->setRange(0, 100);
    m_progressBar->setValue(0);
    mainLayout->addWidget(m_progressBar);

    // 结果表格
    m_tableResults = new QTableWidget(this);
    m_tableResults->setColumnCount(13);
    m_tableResults->setHorizontalHeaderLabels({
        "序号", "左目图像", "右目图像",
        "机械臂X", "机械臂Y", "机械臂Z",
        "检测X", "检测Y", "检测Z",
        "误差X", "误差Y", "误差Z", "总误差"
    });
    m_tableResults->horizontalHeader()->setStretchLastSection(true);
    m_tableResults->setEditTriggers(QAbstractItemView::NoEditTriggers);
    m_tableResults->setSelectionBehavior(QAbstractItemView::SelectRows);
    mainLayout->addWidget(m_tableResults, 2);

    // 日志区域
    QGroupBox* logGroup = new QGroupBox("日志", this);
    QVBoxLayout* logLayout = new QVBoxLayout(logGroup);
    m_logEdit = new QTextEdit(this);
    m_logEdit->setReadOnly(true);
    m_logEdit->setMaximumHeight(150);
    logLayout->addWidget(m_logEdit);
    mainLayout->addWidget(logGroup);
}

void BatchVerifyDialog::onSelectDirectory()
{
    QString dirPath = QFileDialog::getExistingDirectory(
        this, "选择数据目录", m_currentDirectory);

    if (dirPath.isEmpty()) {
        return;
    }

    m_currentDirectory = dirPath;
    m_lblDirectory->setText(dirPath);

    // 扫描目录
    if (scanDirectory(dirPath)) {
        m_btnStartVerify->setEnabled(true);
        appendLog(QString("成功加载 %1 组数据").arg(m_items.size()));
    } else {
        m_btnStartVerify->setEnabled(false);
        appendLog("加载数据失败");
    }
}

bool BatchVerifyDialog::scanDirectory(const QString& dirPath)
{
    m_items.clear();
    m_tableResults->setRowCount(0);

    QDir dir(dirPath);
    if (!dir.exists()) {
        QMessageBox::warning(this, "错误", "目录不存在");
        return false;
    }

    // 查找 images 和 poses 子目录
    QDir imagesDir(dir.filePath("images"));
    QDir posesDir(dir.filePath("poses"));

    if (!imagesDir.exists() || !posesDir.exists()) {
        QMessageBox::warning(this, "错误",
            "目录结构不正确\n请确保包含 images/ 和 poses/ 子目录");
        return false;
    }

    // 查找所有 JSON 文件
    QStringList jsonFiles = posesDir.entryList(QStringList() << "s-*.json", QDir::Files, QDir::Name);

    if (jsonFiles.isEmpty()) {
        QMessageBox::warning(this, "错误", "未找到位姿数据文件 (s-*.json)");
        return false;
    }

    appendLog(QString("找到 %1 个位姿文件").arg(jsonFiles.size()));

    // 遍历每个 JSON 文件
    for (const QString& jsonFile : jsonFiles) {
        // 提取样本 ID (例如 s-1.json -> s-1)
        QString sampleId = QFileInfo(jsonFile).baseName();

        // 构建图像路径
        QString leftImagePath = imagesDir.filePath(sampleId + "_L.png");
        QString rightImagePath = imagesDir.filePath(sampleId + "_R.png");

        // 检查图像是否存在
        if (!QFile::exists(leftImagePath) || !QFile::exists(rightImagePath)) {
            appendLog(QString("警告: 样本 %1 缺少图像文件").arg(sampleId));
            continue;
        }

        // 加载 JSON 文件获取机械臂坐标
        QString jsonPath = posesDir.filePath(jsonFile);
        QFile file(jsonPath);
        if (!file.open(QIODevice::ReadOnly)) {
            appendLog(QString("警告: 无法打开 %1").arg(jsonFile));
            continue;
        }

        QByteArray jsonData = file.readAll();
        file.close();

        QJsonDocument doc = QJsonDocument::fromJson(jsonData);
        if (doc.isNull() || !doc.isObject()) {
            appendLog(QString("警告: %1 不是有效的 JSON 文件").arg(jsonFile));
            continue;
        }

        QJsonObject obj = doc.object();

        // 读取 t_cp (机械臂末端位姿)
        if (!obj.contains("t_cp")) {
            appendLog(QString("警告: %1 缺少 t_cp 字段").arg(jsonFile));
            continue;
        }

        QJsonObject tcpObj = obj["t_cp"].toObject();
        QJsonObject translation = tcpObj["translation_mm"].toObject();
        QJsonObject rotation = tcpObj["rotation_quat"].toObject();

        // 创建验证项
        BatchVerifyItem item;
        item.leftImagePath = leftImagePath;
        item.rightImagePath = rightImagePath;
        item.detected = false;
        item.errorTotal = 0;

        // 读取位置 (mm)
        item.robotX = translation["x"].toDouble();
        item.robotY = translation["y"].toDouble();
        item.robotZ = translation["z"].toDouble();

        // 读取四元数
        double qw = rotation["w"].toDouble();
        double qx = rotation["x"].toDouble();
        double qy = rotation["y"].toDouble();
        double qz = rotation["z"].toDouble();

        // 四元数转欧拉角 (ZYX顺序，单位：度)
        // Roll (X轴旋转)
        double sinr_cosp = 2.0 * (qw * qx + qy * qz);
        double cosr_cosp = 1.0 - 2.0 * (qx * qx + qy * qy);
        item.robotRx = std::atan2(sinr_cosp, cosr_cosp) * 180.0 / M_PI;

        // Pitch (Y轴旋转)
        double sinp = 2.0 * (qw * qy - qz * qx);
        if (std::abs(sinp) >= 1)
            item.robotRy = std::copysign(M_PI / 2, sinp) * 180.0 / M_PI;
        else
            item.robotRy = std::asin(sinp) * 180.0 / M_PI;

        // Yaw (Z轴旋转)
        double siny_cosp = 2.0 * (qw * qz + qx * qy);
        double cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz);
        item.robotRz = std::atan2(siny_cosp, cosy_cosp) * 180.0 / M_PI;

        m_items.push_back(item);
    }

    if (m_items.empty()) {
        QMessageBox::warning(this, "错误", "未找到有效的数据");
        return false;
    }

    updateTable();
    appendLog(QString("成功加载 %1 组数据").arg(m_items.size()));
    return true;
}

bool BatchVerifyDialog::loadRobotCoordinates(const QString& filePath)
{
    // 此函数已不再使用，因为坐标直接从 JSON 读取
    return true;
}

void BatchVerifyDialog::onStartVerify()
{
    if (!m_detector) {
        QMessageBox::warning(this, "错误", "检测器未初始化");
        return;
    }

    m_isVerifying = true;
    m_btnStartVerify->setEnabled(false);
    m_btnStopVerify->setEnabled(true);
    m_btnExport->setEnabled(false);

    appendLog("开始批量验证...");

    // 设置检测参数
    m_detector->SetDetectionFlags(true, true, false);

    int successCount = 0;
    int totalCount = m_items.size();

    for (size_t i = 0; i < m_items.size(); ++i) {
        if (!m_isVerifying) {
            appendLog("验证已停止");
            break;
        }

        appendLog(QString("正在验证第 %1/%2 组数据...").arg(i + 1).arg(totalCount));

        if (detectImagePair(m_items[i])) {
            calculateError(m_items[i]);
            successCount++;
        }

        // 更新进度
        m_progressBar->setValue((i + 1) * 100 / totalCount);
        updateTable();

        // 处理事件，保持界面响应
        QApplication::processEvents();
    }

    m_isVerifying = false;
    m_btnStartVerify->setEnabled(true);
    m_btnStopVerify->setEnabled(false);
    m_btnExport->setEnabled(true);

    appendLog(QString("验证完成: 成功 %1/%2").arg(successCount).arg(totalCount));
}

void BatchVerifyDialog::onStopVerify()
{
    m_isVerifying = false;
    appendLog("正在停止验证...");
}

bool BatchVerifyDialog::detectImagePair(BatchVerifyItem& item)
{
    // 加载左目图像
    QImage leftImage(item.leftImagePath);
    if (leftImage.isNull()) {
        appendLog(QString("错误: 无法加载左目图像 %1").arg(item.leftImagePath));
        return false;
    }

    // 加载右目图像
    QImage rightImage(item.rightImagePath);
    if (rightImage.isNull()) {
        appendLog(QString("错误: 无法加载右目图像 %1").arg(item.rightImagePath));
        return false;
    }

    // 准备相机内参
    CameraIntrinsics intrinsics;
    intrinsics.fx = m_sbFx->value();
    intrinsics.fy = m_sbFy->value();
    intrinsics.cx = m_sbCx->value();
    intrinsics.cy = m_sbCy->value();

    // 检测左目标定板
    QImage leftRgb = leftImage.convertToFormat(QImage::Format_RGB888);
    ChessboardDetectResult leftResult;
    int ret = m_detector->DetectChessboardWithPose(
        leftRgb.bits(),
        leftRgb.width(),
        leftRgb.height(),
        3,
        m_sbPatternWidth->value(),
        m_sbPatternHeight->value(),
        m_sbSquareSize->value(),
        intrinsics,
        leftResult);

    // 检测右目标定板
    QImage rightRgb = rightImage.convertToFormat(QImage::Format_RGB888);
    ChessboardDetectResult rightResult;
    int retRight = m_detector->DetectChessboardWithPose(
        rightRgb.bits(),
        rightRgb.width(),
        rightRgb.height(),
        3,
        m_sbPatternWidth->value(),
        m_sbPatternHeight->value(),
        m_sbSquareSize->value(),
        intrinsics,
        rightResult);

    if (ret == 0 && leftResult.detected && retRight == 0 && rightResult.detected) {
        item.detected = true;

        if (leftResult.hasPose) {
            item.camX = leftResult.center.x;
            item.camY = leftResult.center.y;
            item.camZ = leftResult.center.z;
            item.camRx = leftResult.eulerAngles[0];  // Roll
            item.camRy = leftResult.eulerAngles[1];  // Pitch
            item.camRz = leftResult.eulerAngles[2];  // Yaw
        }

        return true;
    } else {
        item.detected = false;
        appendLog(QString("警告: 检测失败 - %1").arg(QFileInfo(item.leftImagePath).fileName()));
        return false;
    }
}

void BatchVerifyDialog::calculateError(BatchVerifyItem& item)
{
    if (!item.detected) {
        item.errorX = item.errorY = item.errorZ = item.errorTotal = 0;
        return;
    }

    // 计算位置误差
    item.errorX = item.camX - item.robotX;
    item.errorY = item.camY - item.robotY;
    item.errorZ = item.camZ - item.robotZ;
    item.errorTotal = std::sqrt(item.errorX * item.errorX +
                                item.errorY * item.errorY +
                                item.errorZ * item.errorZ);
}

void BatchVerifyDialog::updateTable()
{
    m_tableResults->setRowCount(m_items.size());

    for (size_t i = 0; i < m_items.size(); ++i) {
        const BatchVerifyItem& item = m_items[i];

        m_tableResults->setItem(i, 0, new QTableWidgetItem(QString::number(i + 1)));
        m_tableResults->setItem(i, 1, new QTableWidgetItem(QFileInfo(item.leftImagePath).fileName()));
        m_tableResults->setItem(i, 2, new QTableWidgetItem(QFileInfo(item.rightImagePath).fileName()));

        m_tableResults->setItem(i, 3, new QTableWidgetItem(QString::number(item.robotX, 'f', 3)));
        m_tableResults->setItem(i, 4, new QTableWidgetItem(QString::number(item.robotY, 'f', 3)));
        m_tableResults->setItem(i, 5, new QTableWidgetItem(QString::number(item.robotZ, 'f', 3)));

        if (item.detected) {
            m_tableResults->setItem(i, 6, new QTableWidgetItem(QString::number(item.camX, 'f', 3)));
            m_tableResults->setItem(i, 7, new QTableWidgetItem(QString::number(item.camY, 'f', 3)));
            m_tableResults->setItem(i, 8, new QTableWidgetItem(QString::number(item.camZ, 'f', 3)));

            m_tableResults->setItem(i, 9, new QTableWidgetItem(QString::number(item.errorX, 'f', 3)));
            m_tableResults->setItem(i, 10, new QTableWidgetItem(QString::number(item.errorY, 'f', 3)));
            m_tableResults->setItem(i, 11, new QTableWidgetItem(QString::number(item.errorZ, 'f', 3)));
            m_tableResults->setItem(i, 12, new QTableWidgetItem(QString::number(item.errorTotal, 'f', 3)));
        } else {
            m_tableResults->setItem(i, 6, new QTableWidgetItem("未检测"));
            m_tableResults->setItem(i, 7, new QTableWidgetItem("-"));
            m_tableResults->setItem(i, 8, new QTableWidgetItem("-"));
            m_tableResults->setItem(i, 9, new QTableWidgetItem("-"));
            m_tableResults->setItem(i, 10, new QTableWidgetItem("-"));
            m_tableResults->setItem(i, 11, new QTableWidgetItem("-"));
            m_tableResults->setItem(i, 12, new QTableWidgetItem("-"));
        }
    }

    m_tableResults->resizeColumnsToContents();
}

void BatchVerifyDialog::onExportResults()
{
    QString fileName = QFileDialog::getSaveFileName(
        this, "导出验证结果", "",
        "CSV文件 (*.csv);;文本文件 (*.txt)");

    if (fileName.isEmpty()) {
        return;
    }

    QFile file(fileName);
    if (!file.open(QIODevice::WriteOnly | QIODevice::Text)) {
        QMessageBox::warning(this, "错误", "无法创建文件");
        return;
    }

    QTextStream out(&file);
    out.setCodec("UTF-8");

    // 写入表头
    out << "序号,左目图像,右目图像,机械臂X,机械臂Y,机械臂Z,检测X,检测Y,检测Z,误差X,误差Y,误差Z,总误差\n";

    // 写入数据
    for (size_t i = 0; i < m_items.size(); ++i) {
        const BatchVerifyItem& item = m_items[i];
        out << (i + 1) << ","
            << QFileInfo(item.leftImagePath).fileName() << ","
            << QFileInfo(item.rightImagePath).fileName() << ","
            << item.robotX << "," << item.robotY << "," << item.robotZ << ",";

        if (item.detected) {
            out << item.camX << "," << item.camY << "," << item.camZ << ","
                << item.errorX << "," << item.errorY << "," << item.errorZ << ","
                << item.errorTotal;
        } else {
            out << "未检测,-,-,-,-,-,-";
        }
        out << "\n";
    }

    file.close();
    appendLog(QString("结果已导出到: %1").arg(fileName));
    QMessageBox::information(this, "成功", "验证结果已导出");
}

void BatchVerifyDialog::appendLog(const QString& message)
{
    QString timestamp = QDateTime::currentDateTime().toString("hh:mm:ss");
    m_logEdit->append(QString("[%1] %2").arg(timestamp).arg(message));
}