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

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

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

// Constructor
JFJochViewerImage::JFJochViewerImage(JFJochReader &reader,
                                     QWidget *parent) : reader(reader), QGraphicsView(parent) {
    setDragMode(QGraphicsView::NoDrag); // Disable default drag mode
    setTransformationAnchor(QGraphicsView::AnchorUnderMouse); // Zoom anchors
    setRenderHint(QPainter::Antialiasing); // Enable smooth rendering
    setRenderHint(QPainter::SmoothPixmapTransform);

    setFocusPolicy(Qt::ClickFocus);
    // Connect the horizontal scrollbar's valueChanged signal
    connect(horizontalScrollBar(), &QScrollBar::valueChanged, this, &JFJochViewerImage::onScroll);

    // Connect the vertical scrollbar's valueChanged signal
    connect(verticalScrollBar(), &QScrollBar::valueChanged, this, &JFJochViewerImage::onScroll);
}

// Handle zooming with mouse wheel
void JFJochViewerImage::wheelEvent(QWheelEvent *event) {
    const double zoomFactor = 1.15; // Zoom factor

    // Get the position of the mouse in scene coordinates
    QPointF targetScenePos = mapToScene(event->position().toPoint());

    if (event->modifiers() == Qt::ShiftModifier) {
        float new_foreground = foreground + event->angleDelta().y() / 120.0f;
        if (new_foreground < 1)
            new_foreground = 1.0;
        foreground = new_foreground;
        emit foregroundChanged(foreground);
        Redraw();
    } else {
        // Perform zooming
        if (event->angleDelta().y() > 0) {
            if (scale_factor * zoomFactor < 500.0) {
                scale_factor *= zoomFactor;
                scale(zoomFactor, zoomFactor);
            }
        } else {
            if (scale_factor > 0.2) {
                scale_factor *= 1.0 / zoomFactor;
                scale(1.0 / zoomFactor, 1.0 / zoomFactor);
            }
        }

        // Adjust the view's center to keep the zoom focused on the mouse position
        QPointF updatedViewportCenter = mapToScene(viewport()->rect().center());
        QPointF delta = targetScenePos - updatedViewportCenter;
        translate(delta.x(), delta.y()); // Shift the view

        updateOverlay();
    }
}

// Handle start of panning on mouse press
void JFJochViewerImage::mousePressEvent(QMouseEvent *event) {
    if (event->button() == Qt::LeftButton) {
        setCursor(Qt::ClosedHandCursor); // Change cursor to indicate panning
        lastMousePos = event->pos();
    }
    QGraphicsView::mousePressEvent(event); // Call base implementation
}

// Handle panning while moving the mouse
void JFJochViewerImage::mouseMoveEvent(QMouseEvent *event) {
    if (event->buttons() & Qt::LeftButton) {
        QPointF delta = mapToScene(event->pos()) - mapToScene(lastMousePos);
        lastMousePos = event->pos();
        translate(delta.x(), delta.y()); // Adjust the view's pan
        updateOverlay();
    }
    QGraphicsView::mouseMoveEvent(event);
}


void JFJochViewerImage::LoadImageInternal() {
    if (!image || (image->dataset->image_size_x <= 0))
        return;

    // Create a QImage with RGB format
    QImage qimage(image->dataset->image_size_x, image->dataset->image_size_y, QImage::Format_RGB888);
    image_rgb.resize(image->dataset->image_size_x * image->dataset->image_size_y);

    rgb sat_color{};

    int r,g,b,a;
    feature_color.getRgb(&r, &g, &b, &a);
    auto bad_color = rgb{.r = static_cast<uint8_t>(r), .g = static_cast<uint8_t>(g), .b= static_cast<uint8_t>(b)};

    if (show_saturation) {
        sat_color = bad_color;
    } else
        sat_color = color_scale.Apply(1.0);

    // Fill the QImage with pixel data from the array
    for (int y = 0; y < image->dataset->image_size_y; ++y) {
        uchar *scanLine = qimage.scanLine(y); // Get writable pointer to the row
        for (int x = 0; x < image->dataset->image_size_x; ++x) {
            auto pxl = x + y * image->dataset->image_size_x;

            const double val_orig = (image->image[pxl] - background);
            double val = val_orig / (foreground - background);
            if (image->image[pxl] == GAP_PXL_VALUE)
                image_rgb[pxl] = color_scale.Apply(ColorScaleSpecial::Gap);
            else if (image->image[pxl] == ERROR_PXL_VALUE)
                image_rgb[pxl] = bad_color;
            else if (image->image[pxl] == SATURATED_PXL_VALUE)
                image_rgb[pxl] = sat_color;
            else if (val >= 1.0)
                image_rgb[pxl] = color_scale.Apply(1.0);
            else if (val >= 0.0)
                image_rgb[pxl] = color_scale.Apply(val);
            else
                image_rgb[pxl] = color_scale.Apply(0.0);

            scanLine[x * 3 + 0] = image_rgb[pxl].r; // Red
            scanLine[x * 3 + 1] = image_rgb[pxl].g; // Green
            scanLine[x * 3 + 2] = image_rgb[pxl].b; // Blue
        }
    }

    // Convert QImage to QPixmap and return
    auto pixmap = QPixmap::fromImage(qimage);

    // Display the pixmap in a QGraphicsView
    auto scene = new QGraphicsScene(this);
    auto pixmapItem = new QGraphicsPixmapItem(pixmap);
    scene->addItem(pixmapItem);
    setScene(scene);
}

void JFJochViewerImage::DrawSpots() {
    for (const auto &s: image->spots) {
        size_t spot_size = 3;

        QColor pen_color = spot_color;
        if (s.indexed)
            pen_color = feature_color;

        QPen pen(pen_color, 10);
        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);
        overlay.push_back(rect);
    }
}

void JFJochViewerImage::DrawResolutionRings() {
    // 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->image_size_x), 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->image_size_y), static_cast<int>(std::ceil(visibleRect.bottom())));

    if (res_ring_auto) {
        float radius_x_0 = image->dataset->geom.GetBeamX_pxl() - startX;
        float radius_x_1 = endX - image->dataset->geom.GetBeamX_pxl();

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

        float radius_y_0 = image->dataset->geom.GetBeamY_pxl() - startY;
        float radius_y_1 = endY - image->dataset->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 = {image->dataset->geom.PxlToRes(radius / 2.0f),
                        image->dataset->geom.PxlToRes(radius / 1.02f)};
        else
            res_ring = {};
    }

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

    for (const auto &d: res_ring) {
        float r = image->dataset->geom.ResToPxl(d);
        QRectF boundingRect(image->dataset->geom.GetBeamX_pxl() - r, image->dataset->geom.GetBeamY_pxl() - r, r * 2,r * 2);
        auto circ = scene()->addEllipse(boundingRect, pen);

        overlay.push_back(circ);

        QPointF point_1(image->dataset->geom.GetBeamX_pxl(), image->dataset->geom.GetBeamY_pxl() - r + 3);
        QPointF point_2(image->dataset->geom.GetBeamX_pxl(), image->dataset->geom.GetBeamY_pxl() + r - 3);
        QPointF point_3(image->dataset->geom.GetBeamX_pxl() - r + 3, image->dataset->geom.GetBeamY_pxl());
        QPointF point_4(image->dataset->geom.GetBeamX_pxl() + r - 3, image->dataset->geom.GetBeamY_pxl());
        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", 8);

            QGraphicsTextItem *textItem = scene()->addText(
                    QString("%1 A").arg(QString::number(d, 'g', 2)), font);
            textItem->setDefaultTextColor(feature_color);
            textItem->setPos(point.value());
            overlay.push_back(textItem);
        }
    }
}

void JFJochViewerImage::DrawBeamCenter() {
    DrawCross(image->dataset->geom.GetBeamX_pxl(), image->dataset->geom.GetBeamY_pxl(), 25, 5, 2);
}

void JFJochViewerImage::DrawTopPixels() {
    int i = 0;
    for (auto iter = image->valid_pixel.crbegin();
         iter != image->valid_pixel.rend() && i < show_highest_pixels;
         iter++, i++)
        DrawCross(iter->second % image->dataset->image_size_x + 0.5, iter->second / image->dataset->image_size_x + 0.5, 15, 3);
}

void JFJochViewerImage::updateOverlay() {
    if (!scene() || !image) return;

    for (auto &i: overlay)
        delete i;
    overlay.clear();

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

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

    // Calculate the range of pixels to process
    int startX = std::max(0, static_cast<int>(std::floor(visibleRect.left())));
    int endX = std::min(static_cast<int>(image->dataset->image_size_x), 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->image_size_y), static_cast<int>(std::ceil(visibleRect.bottom())));

    if (scale_factor > 80.0 && (endX - startX + 1) * (endY - startY + 1) < 2000) {
        // Iterate through the pixels within the visible range
        for (int y = startY; y < endY; ++y) {
            for (int x = startX; x < endX; ++x) {
                QString pixelText;
                int32_t val = image->image[x + y * image->dataset->image_size_x];

                if (val == SATURATED_PXL_VALUE)
                    pixelText = "Sat";
                else if (val == GAP_PXL_VALUE)
                    pixelText = "Gap";
                else if (val == ERROR_PXL_VALUE)
                    pixelText = "Err";
                else
                    pixelText = QString("%1").arg(val);

                // Add or update text in the scene
                QGraphicsTextItem *textItem = scene()->addText(pixelText, font);
                if (luminance(image_rgb[x + y * image->dataset->image_size_x]) > 128.0)
                    textItem->setDefaultTextColor(Qt::black); // Text color
                else
                    textItem->setDefaultTextColor(Qt::white); // Text color
                textItem->setPos(x-0.7, y-0.8); // Position the text over the pixel
                textItem->setScale(0.2); // Scale down to 10% of the original size

                overlay.push_back(textItem);
            }
        }
    }

    DrawBeamCenter();
    if (show_spots)
        DrawSpots();
    DrawResolutionRings();
    DrawTopPixels();
    if (show_saturation)
        DrawSaturation();
}

void JFJochViewerImage::resizeEvent(QResizeEvent *event) {
    // Call the base class implementation first
    QGraphicsView::resizeEvent(event);
    updateOverlay();
}

void JFJochViewerImage::onScroll(int value) {
    updateOverlay();
}

void JFJochViewerImage::Redraw() {
    if (!image) return;

    // Save the current transformation (zoom state)
    QTransform currentTransform = this->transform();

    // Regenerate the image
    LoadImageInternal();

    // Restore the zoom level
    this->setTransform(currentTransform, false); // "false" prevents resetting the view
    updateOverlay();
}

void JFJochViewerImage::loadImage() {
    image = reader.CopyImage();
    Redraw();
}

void JFJochViewerImage::changeBackground(float val) {
    background = val;
    Redraw();
}

void JFJochViewerImage::changeForeground(float val) {
    foreground = val;
    Redraw();
}

void JFJochViewerImage::setColorMap(int color_map) {
    try {
        color_scale.Select(static_cast<ColorScaleEnum>(color_map));
        Redraw();
    } catch (...) {}
}

void JFJochViewerImage::noImage() {
    image.reset();
    if (scene())
        scene()->clear();
}

void JFJochViewerImage::setResolutionRing(QVector<float> v) {
    res_ring = v;
    res_ring_auto = false;
    Redraw();
}

void JFJochViewerImage::setResolutionRingAuto() {
    res_ring_auto = true;
    Redraw();
}

void JFJochViewerImage::showSpots(bool input) {
    show_spots = input;
    Redraw();
}

void JFJochViewerImage::setFeatureColor(QColor input) {
    feature_color = input;
    Redraw();
}

void JFJochViewerImage::setSpotColor(QColor input) {
    spot_color = input;
    Redraw();
}

void JFJochViewerImage::showHighestPixels(int32_t v) {
    show_highest_pixels = v;
    Redraw();
}

void JFJochViewerImage::DrawSaturation() {
    for (const auto &iter: image->saturated_pixel)
        DrawCross(iter % image->dataset->image_size_x + 0.5, iter / image->dataset->image_size_y + 0.5, 20, 4);
}

void JFJochViewerImage::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);

    overlay.push_back(horizontalLine);
    overlay.push_back(verticalLine);

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

void JFJochViewerImage::showSaturation(bool input) {
    show_saturation = input;
    Redraw();
}