Jungfraujoch/viewer/windows/JFJochViewer3DDiffractionWindow.cpp

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only


#include "JFJochViewer3DDiffractionWindow.h"

#include <Qt3DRender/QPointLight>
#include <Qt3DRender/QCamera>
#include <Qt3DExtras/QPhongAlphaMaterial>
#include <QVBoxLayout>
#include <QCloseEvent>

using V3 = QVector3D;

static QMatrix4x4 fromBasisPose(const V3& t, const V3& z, const V3& x) {
    // Build a pose matrix from origin t, forward z, right x; y = z x x
    V3 zz = z.normalized();
    V3 xx = x.normalized();
    V3 yy = QVector3D::crossProduct(zz, xx).normalized();
    QMatrix4x4 M;
    M.setColumn(0, QVector4D(xx, 0));
    M.setColumn(1, QVector4D(yy, 0));
    M.setColumn(2, QVector4D(zz, 0));
    M.setColumn(3, QVector4D(t, 1));
    return M;
}

JFJochViewer3DDiffractionWindow::JFJochViewer3DDiffractionWindow(QWidget* parent) : QMainWindow(parent) {
    setWindowTitle("Diffraction 3D");
    auto centralWidget = new QWidget(this);
    setCentralWidget(centralWidget);

    view = new Qt3DExtras::Qt3DWindow();
    container = QWidget::createWindowContainer(view, centralWidget);

    auto layout = new QVBoxLayout(centralWidget);
    layout->setContentsMargins(0,0,0,0);
    layout->addWidget(container);

    root = new Qt3DCore::QEntity();
    view->setRootEntity(root);

    // Camera + controller
    auto cam = view->camera();
    cam->lens()->setPerspectiveProjection(45.f, 1.6f, 0.01f, 1000.f);
    cam->setPosition(V3(3, 2, 3));
    cam->setViewCenter(V3(0, 0, 0));
    auto camCtrl = new Qt3DExtras::QOrbitCameraController(root);
    camCtrl->setCamera(cam);

    // Light
    auto light = new Qt3DRender::QPointLight(root);
    auto lightEntity = new Qt3DCore::QEntity(root);
    lightEntity->addComponent(light);
    auto lt = new Qt3DCore::QTransform();
    lt->setTranslation(V3(5, 5, 5));
    lightEntity->addComponent(lt);

    rebuildScene();
}

void JFJochViewer3DDiffractionWindow::imageLoaded(std::shared_ptr<const JFJochReaderImage> image) {
    if (image) {
        geom_ = image->Dataset().experiment.GetDiffractionGeometry();
        lattice_ = image->ImageData().indexing_lattice;
        axis_ = image->Dataset().experiment.GetGoniometer();
    } else {
        lattice_ = std::nullopt;
        axis_ = std::nullopt;
    }
    rebuildScene();
}

Qt3DCore::QEntity* JFJochViewer3DDiffractionWindow::makeArrow(Qt3DCore::QEntity* parent, const V3& from, const V3& to,
                                                 const QColor& color, float radius) {
    auto group = new Qt3DCore::QEntity(parent);
    V3 dir = to - from;
    float len = dir.length();
    if (len <= 1e-6f) return group;
    V3 z = dir / len;
    V3 x = qAbs(QVector3D::dotProduct(z, V3(0,1,0))) > 0.9f ? V3(1,0,0) : V3(0,1,0);
    QMatrix4x4 basis = fromBasisPose(from, z, x);

    // Shaft
    auto shaft = new Qt3DCore::QEntity(group);
    auto shaftMesh = new Qt3DExtras::QCylinderMesh();
    shaftMesh->setLength(len * 0.85f);
    shaftMesh->setRadius(radius);
    auto shaftTr = new Qt3DCore::QTransform();
    QMatrix4x4 M = basis;
    M.translate(0, 0, len * 0.425f);
    shaftTr->setMatrix(M);
    auto mat = new Qt3DExtras::QPhongMaterial(group);
    mat->setDiffuse(color);
    shaft->addComponent(shaftMesh);
    shaft->addComponent(shaftTr);
    shaft->addComponent(mat);

    // Head
    auto head = new Qt3DCore::QEntity(group);
    auto headMesh = new Qt3DExtras::QConeMesh();
    headMesh->setLength(len * 0.15f);
    headMesh->setTopRadius(0.f);
    headMesh->setBottomRadius(radius * 2.2f);
    auto headTr = new Qt3DCore::QTransform();
    QMatrix4x4 H = basis;
    H.translate(0, 0, len * 0.925f);
    headTr->setMatrix(H);
    head->addComponent(headMesh);
    head->addComponent(headTr);
    head->addComponent(mat);
    return group;
}

Qt3DCore::QEntity* JFJochViewer3DDiffractionWindow::makePlane(Qt3DCore::QEntity* parent, const QSizeF& size,
                                                 const QMatrix4x4& pose, const QColor& color) {
    auto e = new Qt3DCore::QEntity(parent);
    auto mesh = new Qt3DExtras::QPlaneMesh();
    mesh->setWidth(size.width());
    mesh->setHeight(size.height());
    mesh->setMeshResolution(QSize(2,2));
    auto tr = new Qt3DCore::QTransform();
    tr->setMatrix(pose);
    auto mat = new Qt3DExtras::QPhongAlphaMaterial(e);
    mat->setDiffuse(color);
    mat->setAlpha(0.3f);
    e->addComponent(mesh);
    e->addComponent(tr);
    e->addComponent(mat);
    return e;
}

void JFJochViewer3DDiffractionWindow::rebuildScene() {
    // Clear previous children
    for (auto* child : root->children()) child->setParent(nullptr);

    // Sample: sphere at origin
    {
        sampleEntity = new Qt3DCore::QEntity(root);
        auto mesh = new Qt3DExtras::QSphereMesh();
        mesh->setRadius(0.2f);
        auto tr = new Qt3DCore::QTransform();
        tr->setTranslation(V3(0,0,0));
        auto mat = new Qt3DExtras::QPhongAlphaMaterial(sampleEntity);
        mat->setDiffuse(QColor("#6fa8dc"));
        mat->setAlpha(0.6f);
        sampleEntity->addComponent(mesh);
        sampleEntity->addComponent(tr);
        sampleEntity->addComponent(mat);
    }

    // Lattice vectors as arrows (scale to reasonable scene units)
    if (lattice_.has_value()) {
        latticeEntity = new Qt3DCore::QEntity(root);
        auto a = lattice_->Vec0(); auto b = lattice_->Vec1(); auto c = lattice_->Vec2();
        float s = 0.01f; // scale Angstroms to scene units (1 unit = 100 Å)
        auto A = makeArrow(latticeEntity, V3(0,0,0), V3(a.x*s, a.y*s, a.z*s), QColor("#e6194B"));
        auto B = makeArrow(latticeEntity, V3(0,0,0), V3(b.x*s, b.y*s, b.z*s), QColor("#3cb44b"));
        auto C = makeArrow(latticeEntity, V3(0,0,0), V3(c.x*s, c.y*s, c.z*s), QColor("#4363d8"));
        Q_UNUSED(A); Q_UNUSED(B); Q_UNUSED(C);
    }

    // Detector plane using PONI rotation and beam/distance
    {
        // Build detector basis in lab coords.
        // Assume beam travels +Z from sample at origin.
        // Detector center is at z = distance (meters), with X/Y axes rotated by PONI.
        const float px_mm = geom_.GetPixelSize_mm();
        // Example: 4k x 4k display area; replace with your actual detector size if available.
        const int nx = 4096, ny = 4096;
        const float w = nx * px_mm / 1000.f; // meters
        const float h = ny * px_mm / 1000.f; // meters
        const float zc = geom_.GetDetectorDistance_mm() / 1000.f;

        // PONI rotation matrix (assumed aligns detector normal)
        const auto& R = geom_.GetPoniRotMatrix();
        const auto x = R * Coord(1,0,0);
        const auto y = R * Coord(0,1,0);
        const auto z = R * Coord(0,0,1);
        // Convert RotMatrix to QMatrix4x4
        QMatrix4x4 rot;
        rot.setRow(0, QVector4D(x[0], x[1], x[2], 0));
        rot.setRow(1, QVector4D(y[0], y[1], y[2], 0));
        rot.setRow(2, QVector4D(z[0], z[1], z[2], 0));
        rot.setRow(3, QVector4D(0,0,0,1));

        QMatrix4x4 pose = rot;
        pose.translate(0, 0, zc);
        detectorEntity = makePlane(root, QSizeF(w, h), pose, QColor("#ff9900"));

        // Draw beam direction
        makeArrow(root, V3(0,0,0), V3(0,0,zc*1.1f), Qt::yellow, 0.01f);
    }

    // Rotation axis (if present)
    if (axis_.has_value()) {
        axisEntity = new Qt3DCore::QEntity(root);
        auto u = axis_->GetAxis();
        V3 p0(0,0,0);
        V3 p1(u.x, u.y, u.z);
        p1.normalize();
        p1 *= 0.6f;
        makeArrow(axisEntity, p0, p1, QColor("#00ffff"), 0.015f);
    }
}

void JFJochViewer3DDiffractionWindow::closeEvent(QCloseEvent *event) {
    event->accept();
    emit closing();
}

void JFJochViewer3DDiffractionWindow::open() {
    show();
    raise();
    activateWindow();
}