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 <QWidget>
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QPushButton>
#include <QByteArray>

#include <Qt3DCore/QTransform>
#include <Qt3DExtras/QCylinderMesh>
#include <Qt3DExtras/QPhongMaterial>
#include <Qt3DExtras/QForwardRenderer>
#include <Qt3DRender/QCamera>
#include <Qt3DRender/QPointLight>
#include <Qt3DRender/QMaterial>
#include <Qt3DRender/QEffect>
#include <Qt3DRender/QTechnique>
#include <Qt3DRender/QRenderPass>
#include <Qt3DRender/QShaderProgram>
#include <Qt3DRender/QGraphicsApiFilter>
#include <Qt3DRender/QFilterKey>
#include <Qt3DRender/QPointSize>
#include <Qt3DExtras/QPerVertexColorMaterial>
#include "../../common/CrystalLattice.h"

namespace {
    // Custom material: per-vertex colored points with a fixed on-screen size.
    // Needed because Qt6's QPerVertexColorMaterial doesn't let us control point size,
    // and QPointSize must live inside the material's render pass to take effect.
    Qt3DRender::QMaterial *MakePointMaterial(Qt3DCore::QNode *parent, float point_size) {
        auto *material = new Qt3DRender::QMaterial(parent);
        auto *effect = new Qt3DRender::QEffect(material);
        auto *technique = new Qt3DRender::QTechnique(effect);

        technique->graphicsApiFilter()->setApi(Qt3DRender::QGraphicsApiFilter::OpenGL);
        technique->graphicsApiFilter()->setProfile(Qt3DRender::QGraphicsApiFilter::CoreProfile);
        technique->graphicsApiFilter()->setMajorVersion(3);
        technique->graphicsApiFilter()->setMinorVersion(3);

        auto *filterKey = new Qt3DRender::QFilterKey(technique);
        filterKey->setName(QStringLiteral("renderingStyle"));
        filterKey->setValue(QStringLiteral("forward"));
        technique->addFilterKey(filterKey);

        const char *vertexShader = R"(
            #version 330 core
            in vec3 vertexPosition;
            in vec3 vertexColor;
            out vec3 fColor;
            uniform mat4 modelViewProjection;
            void main() {
                fColor = vertexColor;
                gl_Position = modelViewProjection * vec4(vertexPosition, 1.0);
                gl_PointSize = 6.0;
            }
        )";

        const char *fragmentShader = R"(
            #version 330 core
            in vec3 fColor;
            out vec4 outColor;
            void main() { outColor = vec4(fColor, 1.0); }
        )";

        auto *shader = new Qt3DRender::QShaderProgram(material);
        shader->setVertexShaderCode(QByteArray(vertexShader));
        shader->setFragmentShaderCode(QByteArray(fragmentShader));

        auto *pass = new Qt3DRender::QRenderPass(technique);
        pass->setShaderProgram(shader);

        auto *ps = new Qt3DRender::QPointSize(pass);
        ps->setSizeMode(Qt3DRender::QPointSize::Fixed);
        ps->setValue(point_size);
        pass->addRenderState(ps);

        technique->addRenderPass(pass);
        effect->addTechnique(technique);
        material->setEffect(effect);
        return material;
    }

    // Draw a cylinder from origin along the given (already scaled) vector.
    Qt3DCore::QEntity *MakeVector(Qt3DCore::QEntity *parent,
                                  const QVector3D &vec,
                                  const QColor &color,
                                  float thickness) {
        auto *entity = new Qt3DCore::QEntity(parent);

        const float length = vec.length();
        if (length < 1e-6f)
            return entity;

        auto *mesh = new Qt3DExtras::QCylinderMesh();
        mesh->setRadius(thickness);
        mesh->setLength(length);

        auto *material = new Qt3DExtras::QPhongMaterial();
        material->setDiffuse(color);

        auto *transform = new Qt3DCore::QTransform();
        // QCylinderMesh is oriented along +Y by default; rotate +Y to vec direction.
        const QVector3D yAxis(0, 1, 0);
        const QVector3D dir = vec.normalized();
        const QQuaternion rot = QQuaternion::rotationTo(yAxis, dir);
        transform->setRotation(rot);
        transform->setTranslation(vec * 0.5f); // cylinder centered at its midpoint

        entity->addComponent(mesh);
        entity->addComponent(material);
        entity->addComponent(transform);
        return entity;
    }
}

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

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

    // 3D view embedded in a QWidget container
    view = new Qt3DExtras::Qt3DWindow();
    view->defaultFrameGraph()->setClearColor(QColor(20, 20, 25));
    QWidget *container = QWidget::createWindowContainer(view, central);
    container->setMinimumSize(400, 400);
    container->setFocusPolicy(Qt::StrongFocus);
    layout->addWidget(container, 1);

    // Controls
    auto *controls = new QHBoxLayout();
    crystalFrameCheck = new QCheckBox("Crystal frame (angle = 0)", central);
    crystalFrameCheck->setChecked(false);
    crystalFrameCheck->setEnabled(false); // only for rotation data
    showCellCheck = new QCheckBox("Show reciprocal cell", central);
    showCellCheck->setChecked(true);
    accumulateCheck = new QCheckBox("Accumulate", central);
    accumulateCheck->setChecked(false);
    auto *clearButton = new QPushButton("Clear accumulated", central);

    controls->addWidget(crystalFrameCheck);
    controls->addWidget(showCellCheck);
    controls->addWidget(accumulateCheck);
    controls->addStretch(1);
    controls->addWidget(clearButton);
    layout->addLayout(controls);

    setCentralWidget(central);

    connect(crystalFrameCheck, &QCheckBox::toggled, this,
            &JFJochViewerReciprocalSpaceWindow::rebuildSpots);   // frame change → reposition points only
    connect(showCellCheck, &QCheckBox::toggled, this,
            &JFJochViewerReciprocalSpaceWindow::rebuildScene);   // cell vectors only
    connect(clearButton, &QPushButton::clicked, this, [this] {
        spots_.clear();
        rebuildSpots();
    });

    // Root entity + camera
    root = new Qt3DCore::QEntity();

    auto *camera = view->camera();
    camera->lens()->setPerspectiveProjection(45.0f, 1.0f, 0.1f, 10000.0f);
    camera->setPosition(QVector3D(0, 0, 80));
    camera->setViewCenter(QVector3D(0, 0, 0));

    auto *lightEntity = new Qt3DCore::QEntity(root);
    auto *light = new Qt3DRender::QPointLight(lightEntity);
    light->setIntensity(1.0f);
    auto *lightTransform = new Qt3DCore::QTransform(lightEntity);
    lightTransform->setTranslation(QVector3D(0, 0, 200));
    lightEntity->addComponent(light);
    lightEntity->addComponent(lightTransform);

    camController = new Qt3DExtras::QOrbitCameraController(root);
    camController->setCamera(camera);

    // --- Persistent point-cloud entity for spots (one draw call for all spots) ---
    spotEntity = new Qt3DCore::QEntity(root);

    auto *geometry = new Qt3DCore::QGeometry(spotEntity);

    spotPosBuffer = new Qt3DCore::QBuffer(geometry);
    spotColorBuffer = new Qt3DCore::QBuffer(geometry);

    spotPosAttr = new Qt3DCore::QAttribute(geometry);
    spotPosAttr->setName(Qt3DCore::QAttribute::defaultPositionAttributeName());
    spotPosAttr->setVertexBaseType(Qt3DCore::QAttribute::Float);
    spotPosAttr->setVertexSize(3);
    spotPosAttr->setAttributeType(Qt3DCore::QAttribute::VertexAttribute);
    spotPosAttr->setBuffer(spotPosBuffer);
    spotPosAttr->setByteStride(3 * sizeof(float));
    geometry->addAttribute(spotPosAttr);

    spotColorAttr = new Qt3DCore::QAttribute(geometry);
    spotColorAttr->setName(Qt3DCore::QAttribute::defaultColorAttributeName());
    spotColorAttr->setVertexBaseType(Qt3DCore::QAttribute::Float);
    spotColorAttr->setVertexSize(3);
    spotColorAttr->setAttributeType(Qt3DCore::QAttribute::VertexAttribute);
    spotColorAttr->setBuffer(spotColorBuffer);
    spotColorAttr->setByteStride(3 * sizeof(float));
    geometry->addAttribute(spotColorAttr);

    spotRenderer = new Qt3DRender::QGeometryRenderer(spotEntity);
    spotRenderer->setGeometry(geometry);
    spotRenderer->setPrimitiveType(Qt3DRender::QGeometryRenderer::Points);

    // Reuse Qt's working per-vertex-color material, but inject a QPointSize render
    // state into its existing render pass so the points are large enough to see.
    auto *spotMaterial = new Qt3DExtras::QPerVertexColorMaterial(spotEntity);
    if (auto *effect = spotMaterial->effect()) {
        const auto techniques = effect->techniques();
        for (auto *technique : techniques) {
            const auto passes = technique->renderPasses();
            for (auto *pass : passes) {
                auto *ps = new Qt3DRender::QPointSize(pass);
                ps->setSizeMode(Qt3DRender::QPointSize::Fixed);
                ps->setValue(8.0f);
                pass->addRenderState(ps);
            }
        }
    }
    spotEntity->addComponent(spotRenderer);
    spotEntity->addComponent(spotMaterial);

    view->setRootEntity(root);

    rebuildSpots();
    rebuildScene();   // builds axes + cell vectors
}

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::datasetLoaded(std::shared_ptr<const JFJochReaderDataset> in_dataset) {
    // New dataset: clean things (metadata update handled separately, later)
    spots_.clear();
    indexed_lattice_.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);

    rebuildScene();
}

void JFJochViewerReciprocalSpaceWindow::imageLoaded(std::shared_ptr<const JFJochReaderImage> image) {
    if (!accumulateCheck->isChecked()) {
        spots_.clear();
        indexed_lattice_.reset();
    }

    if (!image) {
        rebuildScene();
        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;

    // Goniometer back-rotation to bring the image into the angle = 0 crystal frame,
    // same convention as RotationIndexer / XtalOptimizer.
    std::optional<RotMatrix> back_rot;
    if (axis) {
        const float angle_deg = axis->GetAngle_deg(static_cast<float>(image_number))
                                + axis->GetWedge_deg() / 2.0f;
        back_rot = axis->GetTransformationAngle(angle_deg);
    }

    spots_.reserve(spots_.size() + image->ImageData().spots.size());
    for (const auto &s : image->ImageData().spots) {
        AccumulatedSpot acc;
        acc.recip_lab = s.ReciprocalCoord(geom);
        acc.recip_crystal = back_rot ? (back_rot.value() * acc.recip_lab) : acc.recip_lab;
        acc.indexed = s.indexed;
        acc.ice_ring = s.ice_ring;
        spots_.emplace_back(acc);
    }

    // Soft cap so a long accumulation doesn't grow unbounded
    constexpr size_t max_accumulated = 200000;
    if (spots_.size() > max_accumulated)
        spots_.erase(spots_.begin(),
                     spots_.begin() + static_cast<std::ptrdiff_t>(spots_.size() - max_accumulated));

    // Keep the latest indexed lattice, if available
    if (image->ImageData().indexing_lattice)
        indexed_lattice_ = image->ImageData().indexing_lattice;

    rebuildSpots();
    rebuildScene();   // in case the cell/lattice changed
}

void JFJochViewerReciprocalSpaceWindow::clearScene() {
    if (vectorContent) {
        vectorContent->setParent(static_cast<Qt3DCore::QNode *>(nullptr));
        vectorContent->deleteLater();
        vectorContent = nullptr;
    }
}

void JFJochViewerReciprocalSpaceWindow::rebuildScene() {
    // Only rebuilds the (few) axis + reciprocal-cell vectors.
    clearScene();
    vectorContent = new Qt3DCore::QEntity(root);

    addAxes();
    if (showCellCheck->isChecked() && indexed_lattice_)
        addCellVectors();
}

void JFJochViewerReciprocalSpaceWindow::addCellVectors() {
    if (!indexed_lattice_)
        return;

    const Coord astar = indexed_lattice_->Astar();
    const Coord bstar = indexed_lattice_->Bstar();
    const Coord cstar = indexed_lattice_->Cstar();

    const float thickness = 0.15f;
    MakeVector(vectorContent,
               QVector3D(astar.x, astar.y, astar.z) * scene_scale_,
               Qt::red, thickness);
    MakeVector(vectorContent,
               QVector3D(bstar.x, bstar.y, bstar.z) * scene_scale_,
               Qt::green, thickness);
    MakeVector(vectorContent,
               QVector3D(cstar.x, cstar.y, cstar.z) * scene_scale_,
               Qt::blue, thickness);
}

void JFJochViewerReciprocalSpaceWindow::addAxes() {
    // Faint reference axes (lab frame X/Y/Z) for orientation
    const float len = 10.0f;
    const float thickness = 0.05f;
    MakeVector(vectorContent, QVector3D(len, 0, 0), QColor(120, 60, 60), thickness);
    MakeVector(vectorContent, QVector3D(0, len, 0), QColor(60, 120, 60), thickness);
    MakeVector(vectorContent, QVector3D(0, 0, len), QColor(60, 60, 120), thickness);
}

void JFJochViewerReciprocalSpaceWindow::setSpotColor(QColor input) {
    if (!input.isValid())
        return;
    spot_color = input;
    rebuildSpots();
}

void JFJochViewerReciprocalSpaceWindow::setFeatureColor(QColor input) {
    if (!input.isValid())
        return;
    // Matches JFJochDiffractionImage: indexed spots use the feature color
    indexed_color = input;
    rebuildSpots();
}

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

    QByteArray posData;
    QByteArray colData;
    posData.resize(static_cast<int>(spots_.size() * 3 * sizeof(float)));
    colData.resize(static_cast<int>(spots_.size() * 3 * sizeof(float)));

    auto *pos = reinterpret_cast<float *>(posData.data());
    auto *col = reinterpret_cast<float *>(colData.data());

    for (size_t i = 0; i < spots_.size(); i++) {
        const Coord &c = crystal_frame ? spots_[i].recip_crystal : spots_[i].recip_lab;
        pos[3 * i + 0] = c.x * scene_scale_;
        pos[3 * i + 1] = c.y * scene_scale_;
        pos[3 * i + 2] = c.z * scene_scale_;

        const QColor qc = SpotColorFor(spots_[i].indexed, spots_[i].ice_ring);
        col[3 * i + 0] = static_cast<float>(qc.redF());
        col[3 * i + 1] = static_cast<float>(qc.greenF());
        col[3 * i + 2] = static_cast<float>(qc.blueF());
    }

    spotPosBuffer->setData(posData);
    spotColorBuffer->setData(colData);

    spotPosAttr->setCount(static_cast<uint>(spots_.size()));
    spotColorAttr->setCount(static_cast<uint>(spots_.size()));
    spotRenderer->setVertexCount(static_cast<int>(spots_.size()));
}