// 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 "../../common/JFJochMath.h"
#include "../../common/ROIAzimuthal.h"

#include <QPainterPath>
#include <QBrush>
#include <QGraphicsPixmapItem>
#include <QGraphicsSimpleTextItem>
#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);
        addOverlayItem(rect);
    }
}

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 *ellipse = 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);
        addOverlayItem(ellipse);

        // 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;

            // Use QGraphicsSimpleTextItem for much better performance
            auto *textItem = new QGraphicsSimpleTextItem(label);
            textItem->setFont(font);
            textItem->setBrush(pen_color);
            textItem->setPos(text_x, text_y);
            scene()->addItem(textItem);
            addOverlayItem(textItem);
        }
    }
}

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, PI / 2);
        auto [x3,y3] = geom.ResPhiToPxl(d, PI);
        auto [x4,y4] = geom.ResPhiToPxl(d, 3.0 * 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);
        addOverlayItem(scene()->addEllipse(boundingRect, pen));

        auto [x5,y5] = geom.ResPhiToPxl(d, phi_offset + 0);
        auto [x6,y6] = geom.ResPhiToPxl(d, phi_offset + PI / 2);
        auto [x7,y7] = geom.ResPhiToPxl(d, phi_offset + PI);
        auto [x8,y8] = geom.ResPhiToPxl(d, phi_offset + 3.0 * 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

            auto *textItem = new QGraphicsSimpleTextItem(
                QString("%1 Å").arg(QString::number(d, 'f', 2)));
            textItem->setFont(font);
            textItem->setBrush(ring_color);
            textItem->setPos(point.value());
            scene()->addItem(textItem);
            addOverlayItem(textItem);
        }
        phi_offset += 4.0 / 180.0 * 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();
    DrawROIs();
    DrawTopPixels();

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

    DrawResolutionText();
}

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

    const auto &rois = image->Dataset().experiment.ROI().GetROIDefinition();
    auto geom = image->Dataset().experiment.GetDiffractionGeometry();

    // Distinct colours per ROI; loaded ROIs use solid lines (the interactively
    // drawn scratch ROI keeps its dashed feature_color).
    // TODO: align this palette with the ROI colours in the bottom-panel plots.
    static const QColor palette[] = {Qt::cyan, Qt::yellow, QColor(0xff, 0x57, 0x22),
                                     Qt::green, Qt::magenta, QColor(0x21, 0x96, 0xf3)};
    const int palette_size = sizeof(palette) / sizeof(palette[0]);
    int color_index = 0;

    auto fill_brush = [&](const QColor &c) {
        return show_roi_fill ? QBrush(QColor(c.red(), c.green(), c.blue(), 60)) : QBrush(Qt::NoBrush);
    };

    for (const auto &b : rois.boxes) {
        QColor c = palette[color_index++ % palette_size];
        QPen pen(c, 2); pen.setCosmetic(true);
        addOverlayItem(scene()->addRect(b.GetXMin(), b.GetYMin(), b.GetWidth(), b.GetHeight(), pen, fill_brush(c)));
        AddROILabel(b.GetName(), c, b.GetXMin(), b.GetYMin());
    }

    for (const auto &c_roi : rois.circles) {
        QColor c = palette[color_index++ % palette_size];
        QPen pen(c, 2); pen.setCosmetic(true);
        const float r = c_roi.GetRadius_pxl();
        addOverlayItem(scene()->addEllipse(c_roi.GetX() - r, c_roi.GetY() - r, 2 * r, 2 * r, pen, fill_brush(c)));
        AddROILabel(c_roi.GetName(), c, c_roi.GetX(), c_roi.GetY());
    }

    for (const auto &az : rois.azimuthal)
        DrawAzimuthalROI(az, palette[color_index++ % palette_size], geom);
}

void JFJochDiffractionImage::AddROILabel(const std::string &name, const QColor &color, float px, float py) {
    if (!show_roi_labels)
        return;

    // Just the name; per-ROI statistics are shown in the side-panel ROI list.
    auto *text = scene()->addText(QString::fromStdString(name));
    text->setDefaultTextColor(color);
    text->setFlag(QGraphicsItem::ItemIgnoresTransformations);   // constant on-screen size
    text->setPos(px, py);
    addOverlayItem(text);
}

void JFJochDiffractionImage::DrawAzimuthalROI(const ROIAzimuthal &az, const QColor &color,
                                              const DiffractionGeometry &geom) {
    QPen pen(color, 2); pen.setCosmetic(true);
    QBrush brush = show_roi_fill ? QBrush(QColor(color.red(), color.green(), color.blue(), 60))
                                 : QBrush(Qt::NoBrush);

    const float d_inner = az.GetDMax_A();   // larger d -> smaller radius
    const float d_outer = az.GetDMin_A();
    auto deg2rad = [](float d) { return d * static_cast<float>(PI) / 180.0f; };

    // Sample the boundary through the geometry so the wedge matches the ROI footprint.
    // ResPhiToPxl throws when the resolution is too high for the wavelength; skip such ROIs.
    // move_to_start == true begins a new subpath (no connecting line); false continues
    // the current one (used for the radial edge between a sector's outer and inner arc).
    auto add_arc = [&](QPainterPath &path, float d, float phi_a, float phi_b, int steps, bool move_to_start) -> bool {
        for (int i = 0; i <= steps; i++) {
            float phi = phi_a + (phi_b - phi_a) * static_cast<float>(i) / static_cast<float>(steps);
            try {
                auto [px, py] = geom.ResPhiToPxl(d, phi);
                if (move_to_start && i == 0)
                    path.moveTo(px, py);
                else
                    path.lineTo(px, py);
            } catch (...) { return false; }
        }
        return true;
    };

    QPainterPath path;
    if (az.HasPhi()) {
        float phi0 = deg2rad(az.GetPhiMin_deg());
        float phi1 = deg2rad(az.GetPhiMax_deg());
        if (phi1 < phi0) phi1 += 2.0f * static_cast<float>(PI);   // unwrap the sector
        int steps = std::max(8, static_cast<int>((phi1 - phi0) * 180.0f / static_cast<float>(PI) / 2.0f));
        if (!add_arc(path, d_outer, phi0, phi1, steps, true)) return;    // outer arc
        if (!add_arc(path, d_inner, phi1, phi0, steps, false)) return;   // inner arc; radial edges close it
        path.closeSubpath();
    } else {
        path.setFillRule(Qt::OddEvenFill);                        // annulus: two concentric rings
        const float two_pi = 2.0f * static_cast<float>(PI);
        if (!add_arc(path, d_outer, 0, two_pi, 180, true)) return;
        path.closeSubpath();
        if (!add_arc(path, d_inner, 0, two_pi, 180, true)) return;
        path.closeSubpath();
    }

    addOverlayItem(scene()->addPath(path, pen, brush));

    if (show_roi_labels) {
        try {
            auto [px, py] = geom.ResPhiToPxl(d_outer, az.HasPhi() ? deg2rad(az.GetPhiMin_deg()) : 0.0f);
            AddROILabel(az.GetName(), color, px, py);
        } catch (...) {}
    }
}

void JFJochDiffractionImage::showROILabels(bool input) {
    show_roi_labels = input;
    updateOverlay();
}

void JFJochDiffractionImage::showROIFill(bool input) {
    show_roi_fill = input;
    updateOverlay();
}


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();
        resetScenePointers();
        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

    addOverlayItem(horizontalLine);
    addOverlayItem(verticalLine);
}

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
    // NOTE: hover_resolution_item is NOT tracked in overlay_items_ — it is persistent
    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() {
    // hover_resolution_item is NOT in overlay_items_, so the selective clear won't touch it.
    // However, if scene()->clear() is ever called (e.g. on loadImage(nullptr)),
    // the caller must also set hover_resolution_item = nullptr separately.
}

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);
}