Jungfraujoch/viewer/windows/JFJochViewerReciprocalSpaceWindow.cpp

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

#include "JFJochViewerReciprocalSpaceWindow.h"

#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QSurfaceFormat>
#include <QtMath>

#include <algorithm>
#include <cmath>
#include <limits>
#include <QPushButton>

// ============================================================================
//  Shaders
// ============================================================================

// One shader handles both spots (GL_POINTS) and lines (GL_LINES).
// Per-vertex color comes straight from the VBO; MVP is a single uniform.

static const char *kVertSrc = R"(
    #version 330 core
    layout(location = 0) in vec3 aPos;
    layout(location = 1) in vec4 aColor;
    out vec4 vColor;
    uniform mat4 uMVP;
    void main() {
        vColor       = aColor;
        gl_Position  = uMVP * vec4(aPos, 1.0);
        gl_PointSize = 6.0;
    }
)";

static const char *kFragSrc = R"(
    #version 330 core
    in  vec4 vColor;
    out vec4 fragColor;
    uniform bool uPointShape;
    void main() {
        if (uPointShape) {
            // Discard corners to draw a circle instead of a square point.
            vec2 c = gl_PointCoord - vec2(0.5);
            if (dot(c, c) > 0.25) discard;
        }
        fragColor = vColor;
    }
)";

// ============================================================================
//  ReciprocalSpaceGLView
// ============================================================================

ReciprocalSpaceGLView::ReciprocalSpaceGLView(QWidget *parent)
    : QOpenGLWidget(parent)
    , spotsVBO_(QOpenGLBuffer::VertexBuffer)
    , linesVBO_(QOpenGLBuffer::VertexBuffer)
{
    // Request OpenGL 3.3 Core Profile.
    QSurfaceFormat fmt;
    fmt.setVersion(3, 3);
    fmt.setProfile(QSurfaceFormat::CoreProfile);
    fmt.setDepthBufferSize(24);
    setFormat(fmt);
    setFocusPolicy(Qt::StrongFocus);
}

ReciprocalSpaceGLView::~ReciprocalSpaceGLView() {
    makeCurrent();
    spotsVAO_.destroy();
    spotsVBO_.destroy();
    linesVAO_.destroy();
    linesVBO_.destroy();
    doneCurrent();
}

void ReciprocalSpaceGLView::initializeGL() {
    initializeOpenGLFunctions();

    glClearColor(20.f / 255.f, 20.f / 255.f, 25.f / 255.f, 1.0f);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_PROGRAM_POINT_SIZE);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    shader_.addShaderFromSourceCode(QOpenGLShader::Vertex,   kVertSrc);
    shader_.addShaderFromSourceCode(QOpenGLShader::Fragment, kFragSrc);
    shader_.link();

    // Create empty VAO/VBOs so they are valid before the first data arrives.
    setupVAO(spotsVAO_, spotsVBO_);
    setupVAO(linesVAO_, linesVBO_);

    glReady_ = true;

    uploadBuffer(spotsVBO_, spotsVAO_, pendingSpots_);
    spotsCount_ = static_cast<int>(pendingSpots_.size());

    uploadBuffer(linesVBO_, linesVAO_, pendingLines_);
    linesCount_ = static_cast<int>(pendingLines_.size());
}

void ReciprocalSpaceGLView::resizeGL(int w, int h) {
    const float aspect = float(w) / float(h ? h : 1);
    const float halfHeight = zoom_;
    const float halfWidth  = halfHeight * aspect;

    proj_.setToIdentity();
    proj_.ortho(-halfWidth, halfWidth,
                -halfHeight, halfHeight,
                -10000.0f, 10000.0f);
}

QMatrix4x4 ReciprocalSpaceGLView::currentViewMatrix() const {
    QMatrix4x4 view;
    view.translate(0, 0, -zoom_);
    view.rotate(pitch_, 1, 0, 0);
    view.rotate(yaw_,   0, 1, 0);
    view.translate(-target_);
    return view;
}

QMatrix4x4 ReciprocalSpaceGLView::currentMvpMatrix() const {
    return proj_ * currentViewMatrix();
}

void ReciprocalSpaceGLView::paintGL() {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    const QMatrix4x4 mvp = currentMvpMatrix();

    shader_.bind();
    shader_.setUniformValue("uMVP", mvp);

    // Draw axes + cell vectors.
    // Convention: first 6 line vertices are XYZ reference axes,
    // remaining line vertices are reciprocal-cell vectors.
    if (linesCount_ > 0) {
        shader_.setUniformValue("uPointShape", false);
        linesVAO_.bind();

        const int axesVertexCount = std::min(linesCount_, 6);
        if (axesVertexCount > 0) {
            glLineWidth(1.0f);
            glDrawArrays(GL_LINES, 0, axesVertexCount);
        }

        if (linesCount_ > axesVertexCount) {
            // Note: on some OpenGL core-profile drivers wide lines may be clamped to 1 px.
            glLineWidth(3.0f);
            glDrawArrays(GL_LINES, axesVertexCount, linesCount_ - axesVertexCount);
            glLineWidth(1.0f);
        }

        linesVAO_.release();
    }

    // Draw observed spots + predicted nodes + explicit 000 marker.
    if (spotsCount_ > 0) {
        shader_.setUniformValue("uPointShape", true);
        spotsVAO_.bind();
        glDrawArrays(GL_POINTS, 0, spotsCount_);
        spotsVAO_.release();
    }

    shader_.release();
}

// ---------------------------------------------------------------------------

void ReciprocalSpaceGLView::setSpots(const std::vector<Vertex> &spots) {
    pendingSpots_ = spots;
    spotsCount_ = static_cast<int>(pendingSpots_.size());

    if (glReady_) {
        makeCurrent();
        uploadBuffer(spotsVBO_, spotsVAO_, pendingSpots_);
        doneCurrent();
    }

    update();
}

void ReciprocalSpaceGLView::setLines(const std::vector<Vertex> &lines) {
    pendingLines_ = lines;
    linesCount_ = static_cast<int>(pendingLines_.size());

    if (glReady_) {
        makeCurrent();
        uploadBuffer(linesVBO_, linesVAO_, pendingLines_);
        doneCurrent();
    }

    update();
}

void ReciprocalSpaceGLView::resetTarget() {
    target_ = QVector3D(0.0f, 0.0f, 0.0f);
    update();
}

// ---------------------------------------------------------------------------

void ReciprocalSpaceGLView::setupVAO(QOpenGLVertexArrayObject &vao,
                                      QOpenGLBuffer &vbo) {
    vao.create();
    vao.bind();

    vbo.create();
    vbo.bind();

    // location 0: position xyz
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0,
                          3,
                          GL_FLOAT,
                          GL_FALSE,
                          sizeof(Vertex),
                          reinterpret_cast<void *>(offsetof(Vertex, x)));

    // location 1: color rgba
    glEnableVertexAttribArray(1);
    glVertexAttribPointer(1,
                          4,
                          GL_FLOAT,
                          GL_FALSE,
                          sizeof(Vertex),
                          reinterpret_cast<void *>(offsetof(Vertex, r)));

    vbo.release();
    vao.release();
}

void ReciprocalSpaceGLView::uploadBuffer(QOpenGLBuffer &vbo,
                                          QOpenGLVertexArrayObject &vao,
                                          const std::vector<Vertex> &data) {
    vao.bind();
    vbo.bind();

    if (data.empty()) {
        vbo.allocate(nullptr, 0);
    } else {
        vbo.allocate(data.data(),
                     static_cast<int>(data.size() * sizeof(Vertex)));
    }

    vbo.release();
    vao.release();
}

// ---------------------------------------------------------------------------
// Mouse / wheel -> orbit camera + picking
// ---------------------------------------------------------------------------

void ReciprocalSpaceGLView::mousePressEvent(QMouseEvent *e) {
    lastMousePos_ = e->pos();
}

void ReciprocalSpaceGLView::mouseDoubleClickEvent(QMouseEvent *e) {
    if (e->button() == Qt::LeftButton && emitNearestSpot(e->pos())) {
        e->accept();
        return;
    }

    QOpenGLWidget::mouseDoubleClickEvent(e);
}

void ReciprocalSpaceGLView::mouseMoveEvent(QMouseEvent *e) {
    const QPoint delta = e->pos() - lastMousePos_;
    lastMousePos_ = e->pos();

    if (e->buttons() & Qt::LeftButton) {
        yaw_   += delta.x() * 0.5f;
        pitch_ += delta.y() * 0.5f;
        pitch_  = qBound(-89.0f, pitch_, 89.0f);
        update();
    } else if (e->buttons() & Qt::RightButton) {
        const float panScale = zoom_ * 0.0015f;

        QMatrix4x4 rot;
        rot.rotate(pitch_, 1, 0, 0);
        rot.rotate(yaw_,   0, 1, 0);

        const QVector3D right = rot.inverted().mapVector(QVector3D(1, 0, 0));
        const QVector3D up    = rot.inverted().mapVector(QVector3D(0, 1, 0));

        target_ -= right * float(delta.x()) * panScale;
        target_ += up    * float(delta.y()) * panScale;

        update();
    }
}

void ReciprocalSpaceGLView::wheelEvent(QWheelEvent *e) {
    zoom_ *= (e->angleDelta().y() > 0) ? 0.9f : 1.1f;
    zoom_  = qBound(1.0f, zoom_, 5000.0f);
    resizeGL(width(), height());
    update();
}

bool ReciprocalSpaceGLView::emitNearestSpot(const QPoint &screenPos) {
    if (pendingSpots_.empty() || width() <= 0 || height() <= 0)
        return false;

    const QMatrix4x4 mvp = currentMvpMatrix();

    float bestDist2 = std::numeric_limits<float>::max();
    QPointF bestImagePos;
    bool found = false;

    for (const auto &v : pendingSpots_) {
        if (v.pickable <= 0.5f)
            continue;

        const QVector4D clip = mvp * QVector4D(v.x, v.y, v.z, 1.0f);
        if (clip.w() <= 0.0f)
            continue;

        const QVector3D ndc = clip.toVector3DAffine();
        if (ndc.z() < -1.0f || ndc.z() > 1.0f)
            continue;

        const float sx = (ndc.x() * 0.5f + 0.5f) * float(width());
        const float sy = (0.5f - ndc.y() * 0.5f) * float(height());

        const float dx = sx - float(screenPos.x());
        const float dy = sy - float(screenPos.y());
        const float dist2 = dx * dx + dy * dy;

        if (dist2 < bestDist2) {
            bestDist2 = dist2;
            bestImagePos = QPointF(v.image_x, v.image_y);
            found = true;
        }
    }

    // Roughly 15 px picking radius.
    if (!found || bestDist2 > 15.0f * 15.0f)
        return false;

    emit spotDoubleClicked(bestImagePos);
    return true;
}

// ============================================================================
//  JFJochViewerReciprocalSpaceWindow
// ============================================================================

JFJochViewerReciprocalSpaceWindow::JFJochViewerReciprocalSpaceWindow(QWidget *parent)
    : JFJochHelperWindow(parent)
{
    setWindowTitle("Reciprocal space");
    resize(800, 800);

    auto *central = new QWidget(this);
    auto *layout  = new QVBoxLayout(central);

    glView_ = new ReciprocalSpaceGLView(central);
    glView_->setMinimumSize(400, 400);
    layout->addWidget(glView_, 1);

    auto *controls = new QHBoxLayout();

    crystalFrameCheck = new QCheckBox("Crystal frame (angle = 0)", central);
    crystalFrameCheck->setChecked(false);
    crystalFrameCheck->setEnabled(false);

    showCellCheck = new QCheckBox("Show reciprocal cell", central);
    showCellCheck->setChecked(true);

    showPredictedCheck = new QCheckBox("Show predicted nodes", central);
    showPredictedCheck->setChecked(true);

    auto *centerOriginButton = new QPushButton("Center 000", central);

    controls->addWidget(crystalFrameCheck);
    controls->addWidget(showCellCheck);
    controls->addWidget(showPredictedCheck);
    controls->addStretch(1);
    controls->addWidget(centerOriginButton);

    layout->addLayout(controls);
    setCentralWidget(central);

    connect(crystalFrameCheck, &QCheckBox::toggled,
            this, &JFJochViewerReciprocalSpaceWindow::rebuildGL);
    connect(showCellCheck, &QCheckBox::toggled,
            this, &JFJochViewerReciprocalSpaceWindow::rebuildGL);
    connect(showPredictedCheck, &QCheckBox::toggled,
            this, &JFJochViewerReciprocalSpaceWindow::rebuildGL);
    connect(centerOriginButton, &QPushButton::clicked,
            glView_, &ReciprocalSpaceGLView::resetTarget);

    // Reuse the existing helper-window zoom signal. The main window should connect
    // this to JFJochDiffractionImage::centerOnSpot, like the spot/reflection lists.
    connect(glView_, &ReciprocalSpaceGLView::spotDoubleClicked,
            this, &JFJochHelperWindow::zoom);

    rebuildGL();
}

// ---------------------------------------------------------------------------
// Public slots
// ---------------------------------------------------------------------------

void JFJochViewerReciprocalSpaceWindow::datasetLoaded(
        std::shared_ptr<const JFJochReaderDataset> in_dataset)
{
    spots_.clear();
    reflections_.clear();
    indexed_lattice_.reset();
    current_back_rot_.reset();
    has_rotation_ = false;

    if (in_dataset)
        has_rotation_ = in_dataset->experiment.GetGoniometer().has_value();

    crystalFrameCheck->setEnabled(has_rotation_);
    if (!has_rotation_)
        crystalFrameCheck->setChecked(false);

    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::imageLoaded(std::shared_ptr<const JFJochReaderImage> image) {
    spots_.clear();
    reflections_.clear();
    indexed_lattice_.reset();
    current_back_rot_.reset();

    if (!image) {
        rebuildGL();
        return;
    }

    const auto &dataset        = image->Dataset();
    const auto  geom           = dataset.experiment.GetDiffractionGeometry();
    const auto  axis           = dataset.experiment.GetGoniometer();
    const int64_t image_number = image->ImageData().number;

    if (axis) {
        const float angle_deg = axis->GetAngle_deg(static_cast<float>(image_number))
                                + axis->GetWedge_deg() / 2.0f;
        current_back_rot_ = axis->GetTransformationAngle(angle_deg);
    }

    spots_.reserve(image->ImageData().spots.size());

    for (const auto &s : image->ImageData().spots) {
        CurrentSpot spot;
        spot.recip_lab     = s.ReciprocalCoord(geom);
        spot.recip_crystal = current_back_rot_ ? (*current_back_rot_ * spot.recip_lab) : spot.recip_lab;
        spot.image_x       = s.x;
        spot.image_y       = s.y;
        spot.h             = static_cast<int32_t>(s.h);
        spot.k             = static_cast<int32_t>(s.k);
        spot.l             = static_cast<int32_t>(s.l);
        spot.indexed       = s.indexed;
        spot.ice_ring      = s.ice_ring;
        spots_.emplace_back(spot);
    }

    reflections_.reserve(image->ImageData().reflections.size());
    for (const auto &r : image->ImageData().reflections) {
        reflections_.push_back({
            .h = r.h,
            .k = r.k,
            .l = r.l
        });
    }

    if (image->ImageData().indexing_lattice)
        indexed_lattice_ = image->ImageData().indexing_lattice;

    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::setSpotColor(QColor input) {
    if (!input.isValid())
        return;

    spot_color = input;
    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::setFeatureColor(QColor input) {
    if (!input.isValid())
        return;

    indexed_color = input;
    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::setPredictionColor(QColor input) {
    if (!input.isValid())
        return;

    prediction_color = input;
    rebuildGL();
}

// ---------------------------------------------------------------------------
// Private
// ---------------------------------------------------------------------------

QColor JFJochViewerReciprocalSpaceWindow::spotColorFor(bool indexed,
                                                        bool ice_ring) const {
    if (indexed)
        return indexed_color;
    if (ice_ring)
        return ice_ring_color;
    return spot_color;
}

void JFJochViewerReciprocalSpaceWindow::rebuildGL() {
    const bool crystal_frame = crystalFrameCheck->isChecked() && has_rotation_;

    std::optional<CrystalLattice> plot_lattice;
    if (indexed_lattice_) {
        plot_lattice = indexed_lattice_.value();

        // If spots are brought back to angle zero, bring the lattice back too.
        // After this, q = h*Astar + k*Bstar + l*Cstar is in the same frame as the plotted spots.
        if (crystal_frame && current_back_rot_)
            plot_lattice = plot_lattice->Multiply(*current_back_rot_);
    }

    // --- Spots ---------------------------------------------------------------
    std::vector<ReciprocalSpaceGLView::Vertex> spotVerts;
    spotVerts.reserve(spots_.size() + 1);

    auto toF = [](const QColor &c, float &r, float &g, float &b) {
        r = float(c.redF());
        g = float(c.greenF());
        b = float(c.blueF());
    };

    float sr, sg, sb;
    float ir, ig, ib;
    float xr, xg, xb;
    float pr, pg, pb;

    toF(spot_color,       sr, sg, sb);
    toF(indexed_color,    ir, ig, ib);
    toF(ice_ring_color,   xr, xg, xb);
    toF(prediction_color, pr, pg, pb);

    int hMax = 0;
    int kMax = 0;
    int lMax = 0;

    for (const auto &spot : spots_) {
        const Coord &c = crystal_frame ? spot.recip_crystal : spot.recip_lab;

        float r, g, b;
        if (spot.indexed) {
            r = ir;
            g = ig;
            b = ib;
        } else if (spot.ice_ring) {
            r = xr;
            g = xg;
            b = xb;
        } else {
            r = sr;
            g = sg;
            b = sb;
        }

        spotVerts.push_back({
            c.x * scene_scale_,
            c.y * scene_scale_,
            c.z * scene_scale_,
            r, g, b, 1.0f,
            spot.image_x,
            spot.image_y,
            1.0f
        });

        if (spot.indexed) {
            hMax = std::max(hMax, std::abs(spot.h));
            kMax = std::max(kMax, std::abs(spot.k));
            lMax = std::max(lMax, std::abs(spot.l));
        }
    }

    /*
    // Also include predicted-reflection HKLs in the plotting range.
    for (const auto &r : reflections_) {
        hMax = std::max(hMax, std::abs(r.h));
        kMax = std::max(kMax, std::abs(r.k));
        lMax = std::max(lMax, std::abs(r.l));
    } */

    // Explicit 000 reflection marker.
    spotVerts.push_back({
        0.0f, 0.0f, 0.0f,
        1.0f, 1.0f, 1.0f, 1.0f,
        0.0f, 0.0f,
        0.0f
    });

    // Predicted reciprocal-lattice nodes:
    // q = h*Astar + k*Bstar + l*Cstar, in the currently plotted frame.
    if (showPredictedCheck->isChecked() && plot_lattice && hMax > 0 && kMax > 0 && lMax > 0) {
        const Coord a = plot_lattice->Astar();
        const Coord b = plot_lattice->Bstar();
        const Coord c = plot_lattice->Cstar();

        for (int h = -hMax; h <= hMax; ++h) {
            for (int k = -kMax; k <= kMax; ++k) {
                for (int l = -lMax; l <= lMax; ++l) {
                    if (h == 0 && k == 0 && l == 0)
                        continue;

                    bool observedIndexed = false;
                    for (const auto &spot : spots_) {
                        if (spot.indexed && spot.h == h && spot.k == k && spot.l == l) {
                            observedIndexed = true;
                            break;
                        }
                    }

                    if (observedIndexed)
                        continue;

                    bool predictedForImage = false;
                    for (const auto &r : reflections_) {
                        if (r.h == h && r.k == k && r.l == l) {
                            predictedForImage = true;
                            break;
                        }
                    }

                    const Coord node{
                        h * a.x + k * b.x + l * c.x,
                        h * a.y + k * b.y + l * c.y,
                        h * a.z + k * b.z + l * c.z
                    };

                    if (predictedForImage) {
                        // Highlight nodes that are actual predictions for this image.
                        spotVerts.push_back({
                            node.x * scene_scale_,
                            node.y * scene_scale_,
                            node.z * scene_scale_,
                            pr, pg, pb, 0.85f,
                            0.0f, 0.0f,
                            0.0f
                        });
                    }
                }
            }
        }
    }

    // --- Lines: axes + optional cell vectors ---------------------------------
    std::vector<ReciprocalSpaceGLView::Vertex> lineVerts;

    auto addLine = [&](QVector3D a, QVector3D b, QColor col, float alpha = 1.0f) {
        const float r  = float(col.redF());
        const float g  = float(col.greenF());
        const float bv = float(col.blueF());

        lineVerts.push_back({
            a.x(), a.y(), a.z(),
            r, g, bv, alpha,
            0.0f, 0.0f,
            0.0f
        });

        lineVerts.push_back({
            b.x(), b.y(), b.z(),
            r, g, bv, alpha,
            0.0f, 0.0f,
            0.0f
        });
    };

    // Reference coordinate axes: short and white.
    const float len = 20.0f;
    addLine({0, 0, 0}, {len, 0, 0}, Qt::white, 0.75f);
    addLine({0, 0, 0}, {0, len, 0}, Qt::white, 0.75f);
    addLine({0, 0, 0}, {0, 0, len}, Qt::white, 0.75f);

    // Reciprocal cell vectors, using the same plotted lattice frame as predicted nodes.
    if (showCellCheck->isChecked() && plot_lattice) {
        auto addCell = [&](const Coord &v, QColor col) {
            addLine({0, 0, 0},
                    {
                        v.x * scene_scale_,
                        v.y * scene_scale_,
                        v.z * scene_scale_
                    },
                    col,
                    1.0f);
        };

        addCell(plot_lattice->Astar(), Qt::red);
        addCell(plot_lattice->Bstar(), Qt::green);
        addCell(plot_lattice->Cstar(), Qt::blue);
    }

    glView_->setSpots(spotVerts);
    glView_->setLines(lineVerts);
}