Jungfraujoch/viewer/windows/JFJochViewerReciprocalSpaceWindow.h

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

#pragma once

#include <QMainWindow>
#include <QCheckBox>
#include <QOpenGLWidget>
#include <QOpenGLFunctions_3_3_Core>
#include <QOpenGLShaderProgram>
#include <QOpenGLBuffer>
#include <QOpenGLVertexArrayObject>
#include <QMatrix4x4>
#include <QComboBox>
#include <QMouseEvent>
#include <QWheelEvent>

#include <vector>
#include <memory>
#include <optional>

#include "JFJochHelperWindow.h"           // your existing base class
#include "../../common/CrystalLattice.h"  // Coord, RotMatrix, etc.
#include "../../common/SpotToSave.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/GoniometerAxis.h"

// Forward declarations from your codebase
struct JFJochReaderDataset;
class JFJochReaderImage;
struct JFJochReaderSpots;

// ---------------------------------------------------------------------------
// ReciprocalSpaceGLView  –  the actual OpenGL viewport
// ---------------------------------------------------------------------------
class ReciprocalSpaceGLView : public QOpenGLWidget,
                               protected QOpenGLFunctions_3_3_Core {
    Q_OBJECT
public:
    explicit ReciprocalSpaceGLView(QWidget *parent = nullptr);
    ~ReciprocalSpaceGLView() override;

    struct Vertex {
        float x, y, z;
        float r, g, b, a;
        float image_x = 0.0f;
        float image_y = 0.0f;
        float image_number = -1.0f;
        float pickable = 0.0f;
    };

    void resetTarget();

    // Called by the outer window whenever data changes.
    void setSpots(const std::vector<Vertex> &spots);
    void setLines(const std::vector<Vertex> &lines);   // axes + cell vectors

signals:
    void spotDoubleClicked(QPointF imagePos);

protected:
    void initializeGL() override;
    void resizeGL(int w, int h) override;
    void paintGL() override;

    void mousePressEvent(QMouseEvent *e) override;
    void mouseDoubleClickEvent(QMouseEvent *e) override;
    void mouseMoveEvent(QMouseEvent *e) override;
    void wheelEvent(QWheelEvent *e) override;

private:
    void uploadBuffer(QOpenGLBuffer &vbo, QOpenGLVertexArrayObject &vao,
                      const std::vector<Vertex> &data);
    void setupVAO(QOpenGLVertexArrayObject &vao, QOpenGLBuffer &vbo);

    QMatrix4x4 currentViewMatrix() const;
    QMatrix4x4 currentMvpMatrix() const;
    bool emitNearestSpot(const QPoint &screenPos);

    QOpenGLShaderProgram shader_;
    bool                 glReady_ = false;

    QOpenGLVertexArrayObject spotsVAO_;
    QOpenGLBuffer            spotsVBO_;
    int                      spotsCount_ = 0;
    std::vector<Vertex>      pendingSpots_;

    QOpenGLVertexArrayObject linesVAO_;
    QOpenGLBuffer            linesVBO_;
    int                      linesCount_ = 0;
    std::vector<Vertex>      pendingLines_;

    // Camera state
    QMatrix4x4  proj_;
    float       yaw_   = 0.0f;   // degrees
    float       pitch_ = 20.0f;  // degrees
    float       zoom_  = 80.0f;  // distance from camera target
    QVector3D   target_{0.0f, 0.0f, 0.0f};
    QPoint      lastMousePos_;
};

// ---------------------------------------------------------------------------
// JFJochViewerReciprocalSpaceWindow  –  the outer QMainWindow
// ---------------------------------------------------------------------------

class JFJochViewerReciprocalSpaceWindow : public JFJochHelperWindow {
    Q_OBJECT
public:
    explicit JFJochViewerReciprocalSpaceWindow(QWidget *parent = nullptr);

signals:
    void loadSpotsRequest(int64_t start_image, int64_t end_image, int64_t stride);

public slots:
    void datasetLoaded(std::shared_ptr<const JFJochReaderDataset> dataset);
    void imageLoaded(std::shared_ptr<const JFJochReaderImage> image);
    void spotsLoaded(std::shared_ptr<const JFJochReaderSpots> spots);
    void setSpotColor(QColor color);
    void setFeatureColor(QColor color);
    void setPredictionColor(QColor color);

private:
    struct CurrentSpot {
        Coord recip_lab;
        Coord recip_crystal;
        float image_x = 0.0f;
        float image_y = 0.0f;
        int64_t image_number = -1;
        int32_t h = 0;
        int32_t k = 0;
        int32_t l = 0;
        bool  indexed  = false;
        bool  ice_ring = false;
    };

    struct CurrentReflection {
        int32_t h = 0;
        int32_t k = 0;
        int32_t l = 0;
        float image_x = 0.0f;
        float image_y = 0.0f;
        int64_t image_number = -1;
    };

    QColor spotColorFor(bool indexed, bool ice_ring) const;
    void   rebuildGL();   // rebuilds both spot and line vertex data
    void   loadCurrentImageSpots(std::shared_ptr<const JFJochReaderImage> image);
    void   addSpot(const SpotToSave &s,
                   const DiffractionGeometry &geom,
                   const std::optional<GoniometerAxis> &axis,
                   int64_t image_number);

    // UI
    ReciprocalSpaceGLView *glView_            = nullptr;
    QCheckBox             *crystalFrameCheck  = nullptr;
    QCheckBox             *showCellCheck      = nullptr;
    QCheckBox             *showPredictedCheck = nullptr;
    QComboBox             *strideCombo        = nullptr;

    // Data
    std::shared_ptr<const JFJochReaderDataset> current_dataset_;
    std::shared_ptr<const JFJochReaderImage>   current_image_;

    std::vector<CurrentSpot>        spots_;
    std::vector<CurrentReflection>  reflections_;
    std::optional<CrystalLattice>   indexed_lattice_;
    std::optional<RotMatrix>        current_back_rot_;
    bool   has_rotation_ = false;
    bool   full_dataset_mode_ = false;
    float  scene_scale_  = 100.0f;

    // Colors
    QColor spot_color = Qt::green;
    QColor indexed_color = Qt::magenta;   // feature_color in image
    QColor prediction_color = Qt::darkRed;
    QColor ice_ring_color = Qt::cyan;
};