Jungfraujoch/viewer/image_viewer/JFJochDiffractionImage.cpp

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

#include "JFJochDiffractionImage.h"
#include "../../common/DiffractionGeometry.h"

#include <QGraphicsPixmapItem>
#include <QGraphicsScene>
#include <QWheelEvent>
#include <QScrollBar>
#include <QMenu>
#include <cmath>
#include <QMouseEvent>

#include "JFJochSimpleImage.h"

// Constructor
JFJochDiffractionImage::JFJochDiffractionImage(QWidget *parent) : JFJochImage(parent) {}

void JFJochDiffractionImage::mouseHover(QMouseEvent *event) {
    auto coord = mapToScene(event->pos());

    if (image && (coord.x() >= 0)
        && (coord.x() < image->Dataset().experiment.GetXPixelsNum())
        && (coord.y() >= 0)
        && (coord.y() < image->Dataset().experiment.GetYPixelsNum())) {
        float res = image->Dataset().experiment.GetDiffractionGeometry().PxlToRes(coord.x(), coord.y());

        int32_t intensity = image->Image()[std::floor(coord.x()) +
                                           std::floor(coord.y()) * image->Dataset().experiment.GetXPixelsNum()];

        QString intensity_str = QString("I=%1").arg(intensity, 9);

        if (intensity == SATURATED_PXL_VALUE)
            intensity_str = "I=Saturated";
        else if (intensity == GAP_PXL_VALUE)
            intensity_str = "    Gap    ";
        else if (intensity == ERROR_PXL_VALUE)
            intensity_str = "  Bad pxl  ";

        emit writeStatusBar(QString("x=%1 y=%2 %3 d=%4 Å")
            .arg(coord.x(), 0, 'f', 1)
            .arg(coord.y(), 0, 'f', 1)
            .arg(intensity_str)
            .arg(res, 0, 'f', 2));

        // Update hovered resolution text without rebuilding the whole overlay
        hover_resolution = res;
        DrawResolutionText();
        } else {
            emit writeStatusBar("");

            // Clear hover resolution text when outside image
            if (std::isfinite(hover_resolution)) {
                hover_resolution = NAN;
                DrawResolutionText();
            }
        }
}


void JFJochDiffractionImage::LoadImageInternal() {
    if (!image)
        return;

    W = image->Dataset().experiment.GetXPixelsNum();
    H = image->Dataset().experiment.GetYPixelsNum();

    image_fp.resize(W*H);

    auto img = image->Image();
    // Fill the QImage with pixel data from the array
    for (int pxl = 0; pxl < W * H; pxl++) {
        auto val = img[pxl];
        if (val == GAP_PXL_VALUE)
            image_fp[pxl] = NAN;
        else if (val == ERROR_PXL_VALUE)
            image_fp[pxl] = -INFINITY;
        else if (val == SATURATED_PXL_VALUE)
            image_fp[pxl] = INFINITY;
        else
            image_fp[pxl] = static_cast<float>(val);
    }
}

void JFJochDiffractionImage::DrawSpots() {
    // Compute current visible area in scene coordinates
    const QRectF visibleRect = mapToScene(viewport()->geometry()).boundingRect();

    for (const auto &s: image->ImageData().spots) {
        // Skip reflections outside the viewport
        if (!visibleRect.contains(QPointF{s.x, s.y}))
            continue;

        const qreal desired_half_px = 8.0;
        const qreal spot_size = desired_half_px / std::sqrt(std::max(0.0001, scale_factor));

        QColor pen_color = spot_color;
        if (s.indexed)
            pen_color = feature_color;
        else if (highlight_ice_rings && s.ice_ring)
            pen_color = ice_ring_color;

        QPen pen(pen_color, 3);
        pen.setCosmetic(true);

        auto rect = scene()->addRect(s.x - spot_size + 0.5,
                                     s.y - spot_size + 0.5,
                                     2 * spot_size,
                                     2 * spot_size,
                                     pen);
    }
}

void JFJochDiffractionImage::DrawPredictions() {
    QFont font("Arial", 2); // Font for pixel value text
    font.setPixelSize(2); // This will render very small text (1-pixel high).

    const qreal desired_half_px = 8.0;
    const qreal spot_size = desired_half_px / std::sqrt(std::max(0.0001, scale_factor));

    QColor pen_color = prediction_color;

    QPen pen(pen_color, 3);
    pen.setCosmetic(true);

    // Compute current visible area in scene coordinates
    const QRectF visibleRect = mapToScene(viewport()->geometry()).boundingRect();

    for (const auto &s: image->ImageData().reflections) {
        // Skip reflections outside the viewport
        if (!visibleRect.contains(QPointF{s.predicted_x, s.predicted_y}))
            continue;

        auto rect = scene()->addEllipse(s.predicted_x - spot_size + 0.5f,
                                        s.predicted_y - spot_size + 0.5f,
                                        2.0f * spot_size,
                                        2.0f * spot_size,
                                        pen);

        // When zoomed in enough, draw "h k l" above the box
        if (scale_factor >= 10.0) {
            // Format label
            QString label = QString("%1, %2, %3").arg(s.h).arg(s.k).arg(s.l);

            // Position slightly above the top side of the box
            const qreal text_x = s.predicted_x - 5.5f;
            const qreal text_y = s.predicted_y - 10.0f;

            // Add or update text in the scene
            QGraphicsTextItem *textItem = scene()->addText(label, font);
            textItem->setDefaultTextColor(pen_color);
            textItem->setPos(text_x, text_y); // Position the text over the pixel
            // textItem->setScale(1.0); // Scale down to 10% of the original size
        }
    }
}

void JFJochDiffractionImage::DrawResolutionRings() {
    if (ring_mode == RingMode::None)
        return;

    // Get the visible area in the scene coordinates
    QRectF visibleRect = mapToScene(viewport()->geometry()).boundingRect();

    int startX = std::max(0, static_cast<int>(std::floor(visibleRect.left())));
    int endX = std::min(static_cast<int>(image->Dataset().experiment.GetXPixelsNum()),
                        static_cast<int>(std::ceil(visibleRect.right())));
    int startY = std::max(0, static_cast<int>(std::floor(visibleRect.top())));
    int endY = std::min(static_cast<int>(image->Dataset().experiment.GetYPixelsNum()),
                        static_cast<int>(std::ceil(visibleRect.bottom())));

    auto geom = image->Dataset().experiment.GetDiffractionGeometry();
    geom.PoniRot3_rad(0.0);

    QColor ring_color = feature_color;
    if (ring_mode == RingMode::IceRings) {
        ring_color = ice_ring_color;
        res_ring = QVector<float>{ICE_RING_RES_A.begin(), ICE_RING_RES_A.end()};
    } else if (ring_mode == RingMode::Auto) {
        float radius_x_0 = geom.GetBeamX_pxl() - startX;
        float radius_x_1 = endX - geom.GetBeamX_pxl();

        float radius_x = std::max(radius_x_0, radius_x_1);

        float radius_y_0 = geom.GetBeamY_pxl() - startY;
        float radius_y_1 = endY - geom.GetBeamY_pxl();

        float radius_y = std::max(radius_y_0, radius_y_1);

        float radius = std::min(radius_x, radius_y);
        if (radius_x <= 0)
            radius = radius_y;
        if (radius_y <= 0)
            radius = radius_x;

        if (radius > 0)
            res_ring = {
                geom.PxlToRes(radius / 2.0f),
                geom.PxlToRes(radius / 1.02f)
            };
        else
            res_ring = {};
    } else if (ring_mode == RingMode::Estimation) {
        if (image
            && image->ImageData().resolution_estimate
            && std::isfinite(image->ImageData().resolution_estimate.value())
            && image->ImageData().resolution_estimate.value() > 0.0)
            res_ring = {*image->ImageData().resolution_estimate};
        else
            res_ring = {};
    }

    if (res_ring.empty())
        return;

    QPen pen(ring_color, 5);
    pen.setCosmetic(true);

    QVector<qreal> dashPattern = {10, 15};
    pen.setDashPattern(dashPattern);

    float phi_offset = 0;

    float res1 = geom.PxlToRes(0,0);
    float res2 = geom.PxlToRes(image->Dataset().experiment.GetXPixelsNum(),0);
    float res3 = geom.PxlToRes(image->Dataset().experiment.GetXPixelsNum(),image->Dataset().experiment.GetYPixelsNum());
    float res4 = geom.PxlToRes(0,image->Dataset().experiment.GetYPixelsNum());

    float min_res = std::min({res1, res2, res3, res4});

    for (const auto &d: res_ring) {
        if (d < min_res)
            continue;

        auto [x1,y1] = geom.ResPhiToPxl(d, 0);
        auto [x2,y2] = geom.ResPhiToPxl(d, M_PI_2);
        auto [x3,y3] = geom.ResPhiToPxl(d, M_PI);
        auto [x4,y4] = geom.ResPhiToPxl(d, 3.0 * M_PI_2);

        auto x_min = std::min({x1, x2, x3, x4});
        auto x_max = std::max({x1, x2, x3, x4});
        auto y_min = std::min({y1, y2, y3, y4});
        auto y_max = std::max({y1, y2, y3, y4});

        QRectF boundingRect(x_min, y_min, x_max - x_min, y_max - y_min);
        scene()->addEllipse(boundingRect, pen);

        auto [x5,y5] = geom.ResPhiToPxl(d, phi_offset + 0);
        auto [x6,y6] = geom.ResPhiToPxl(d, phi_offset + M_PI_2);
        auto [x7,y7] = geom.ResPhiToPxl(d, phi_offset + M_PI);
        auto [x8,y8] = geom.ResPhiToPxl(d, phi_offset + 3.0 * M_PI_2);

        QPointF point_1(x5, y5);
        QPointF point_2(x6, y6);
        QPointF point_3(x7, y7);
        QPointF point_4(x8, y8);
        std::optional<QPointF> point;

        if (visibleRect.contains(point_1))
            point = point_1;
        else if (visibleRect.contains(point_2))
            point = point_2;
        else if (visibleRect.contains(point_3))
            point = point_3;
        else if (visibleRect.contains(point_4))
            point = point_4;

        if (point) {
            QFont font("Arial", 16);

            const qreal f = std::clamp(scale_factor, 0.5, 50.0);
            font.setPointSizeF(16.0 / sqrt(f)); // base 12pt around scale_factor ~10

            QGraphicsTextItem *textItem = scene()->addText(
                QString("%1 Å").arg(QString::number(d, 'f', 2)), font);
            textItem->setDefaultTextColor(ring_color);
            textItem->setPos(point.value());
        }
        phi_offset += 4.0 / 180.0 * M_PI;
    }
}

void JFJochDiffractionImage::DrawBeamCenter() {
    auto geom = image->Dataset().experiment.GetDiffractionGeometry();

    auto [beam_x, beam_y] = geom.GetDirectBeam_pxl();
    DrawCross(beam_x, beam_y, 25, 5, 2);
}

void JFJochDiffractionImage::DrawTopPixels() {
    int i = 0;
    for (const auto& p : image->GetTopPixels()) {
        if (i >= show_highest_pixels)
            break;

        const int32_t idx = p.second;
        DrawCross(idx % image->Dataset().experiment.GetXPixelsNum() + 0.5,
                  idx / image->Dataset().experiment.GetXPixelsNum() + 0.5, 15, 3);
        i++;
    }
}


void JFJochDiffractionImage::addCustomOverlay() {
    DrawResolutionRings();
    DrawTopPixels();

    DrawBeamCenter();
    if (show_spots)
        DrawSpots();
    if (show_predictions)
        DrawPredictions();
    if (show_saturation)
        DrawSaturation();

    DrawResolutionText();
}


void JFJochDiffractionImage::UpdateForeground() {
    if (!image || !auto_fg)
        return;
    if (hdr_mode) {
        const auto val_range = image->ValidMinMax();
        if (val_range.has_value())
            foreground = val_range->second;
    } else {
        foreground = image->GetAutoContrastValue();
    }
    emit foregroundChanged(foreground);
}

void JFJochDiffractionImage::setHDRMode(bool input) {
    hdr_mode = input;
    UpdateForeground();
    GeneratePixmap();
    Redraw();
}

void JFJochDiffractionImage::loadImage(std::shared_ptr<const JFJochReaderImage> in_image) {
    if (in_image) {
        image = in_image;
        UpdateForeground();
        LoadImageInternal();
        GeneratePixmap();
        Redraw();
        CalcROI();
    } else {
        image.reset();
        W = 0; H = 0;
        if (scene())
            scene()->clear();
        hover_resolution = NAN;
        hover_resolution_item = nullptr;

        CalcROI();
    }
}

void JFJochDiffractionImage::setAutoForeground(bool input) {
    auto_fg = input;
    // If auto_foreground is not set, then view stays with the current settings till these are explicitly changed
    UpdateForeground();
    GeneratePixmap();
    Redraw();
    emit autoForegroundChanged(auto_fg);
}

void JFJochDiffractionImage::setResolutionRing(QVector<float> v) {
    res_ring = v;
    ring_mode = RingMode::Manual;
    updateOverlay();
}

void JFJochDiffractionImage::showSpots(bool input) {
    show_spots = input;
    updateOverlay();
}

void JFJochDiffractionImage::showPredictions(bool input) {
    show_predictions = input;
    updateOverlay();
}

void JFJochDiffractionImage::setSpotColor(QColor input) {
    spot_color = input;
    updateOverlay();
}

void JFJochDiffractionImage::setPredictionColor(QColor input) {
    prediction_color = input;
    updateOverlay();
}

void JFJochDiffractionImage::showHighestPixels(int32_t v) {
    show_highest_pixels = v;
    updateOverlay();
}

void JFJochDiffractionImage::DrawSaturation() {
    for (const auto &iter: image->SaturatedPixels())
        DrawCross(iter % image->Dataset().experiment.GetXPixelsNum() + 0.5,
                  iter / image->Dataset().experiment.GetXPixelsNum() + 0.5, 20, 4);
}

void JFJochDiffractionImage::DrawCross(float x, float y, float size, float width, float z) {
    float sc_size = size / sqrt(scale_factor);

    QPen pen(feature_color, width);
    pen.setCosmetic(true);

    QGraphicsLineItem *horizontalLine = scene()->addLine(x - sc_size, y, x + sc_size, y, pen);
    QGraphicsLineItem *verticalLine = scene()->addLine(x, y - sc_size, x, y + sc_size, pen);

    horizontalLine->setZValue(z); // Ensure it appears above other items
    verticalLine->setZValue(z); // Ensure it appears above other items
}

void JFJochDiffractionImage::showSaturation(bool input) {
    show_saturation = input;
    GeneratePixmap();
    updateOverlay();
}

void JFJochDiffractionImage::highlightIceRings(bool input) {
    highlight_ice_rings = input;
    updateOverlay();
}

void JFJochDiffractionImage::setResolutionRingMode(RingMode mode) {
    ring_mode = mode;
    updateOverlay();
}

void JFJochDiffractionImage::DrawResolutionText() {
    auto scn = scene();
    if (!scn) {
        hover_resolution_item = nullptr; // scene gone
        return;
    }

    // Hide item if no valid hover resolution
    if (!image || !std::isfinite(hover_resolution) || hover_resolution <= 0.0f) {
        if (hover_resolution_item)
            hover_resolution_item->setVisible(false);
        return;
    }

    const QRectF visibleRect = mapToScene(viewport()->geometry()).boundingRect();

    // Fixed on-screen font size (no dependence on scale_factor)
    QFont font("Arial");
    font.setPixelSize(32);  // big, constant size on screen

    const QString label =
        QString("d = %1 Å").arg(QString::number(hover_resolution, 'f', 2));

    // Create the item if it does not exist yet; otherwise reuse it
    if (!hover_resolution_item) {
        hover_resolution_item = scn->addText(label, font);
        hover_resolution_item->setZValue(10.0);
        // Make the text ignore zooming / view transforms
        hover_resolution_item->setFlag(QGraphicsItem::ItemIgnoresTransformations, true);
    } else {
        hover_resolution_item->setFont(font);
        hover_resolution_item->setPlainText(label);
    }

    hover_resolution_item->setDefaultTextColor(feature_color);

    // Keep a roughly constant ~10 px margin by compensating with scale_factor
    const qreal margin_px = 10.0;
    const qreal margin_scene = margin_px / std::max(0.0001, scale_factor);

    QPointF topLeft(visibleRect.left() + margin_scene,
                    visibleRect.top() + margin_scene);
    hover_resolution_item->setPos(topLeft);
    hover_resolution_item->setVisible(true);
}

void JFJochDiffractionImage::beforeOverlayCleared() {
    // The scene is about to clear (and delete) all its items.
    // Drop our non-owning pointer so we never touch a deleted item.
    hover_resolution_item = nullptr;
}

void JFJochDiffractionImage::leaveEvent(QEvent *event) {
    // Mouse left the view: clear hover resolution and hide text
    if (std::isfinite(hover_resolution)) {
        hover_resolution = NAN;
        DrawResolutionText();
    }
    JFJochImage::leaveEvent(event);
}