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

#include "JFJochImage.h"

#include <QGraphicsPixmapItem>
#include <QScrollBar>
#include <QWheelEvent>
#include <QMouseEvent>

JFJochImage::JFJochImage(QWidget *parent) : 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, &JFJochImage::onScroll);

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

    // Optional: a sensible default colormap
    color_scale.Select(ColorScaleEnum::Indigo);
}

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

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

void JFJochImage::changeForeground(float val) {
    foreground = val;
    // Regenerate the image
    GeneratePixmap();
    Redraw();
}

void JFJochImage::setColorMap(int color_map) {
    try {
        color_scale.Select(static_cast<ColorScaleEnum>(color_map));
        // Regenerate the image
        GeneratePixmap();
        Redraw();
    } catch (...) {
    }
}

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

void JFJochImage::wheelEvent(QWheelEvent *event) {
    if (!scene()) return;

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

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

void JFJochImage::resizeEvent(QResizeEvent *event) {
    QGraphicsView::resizeEvent(event);
    updateOverlay();
}

QPointF JFJochImage::RoundPoint(const QPointF &input) {
    return QPointF(qRound(input.x()), qRound(input.y()));
}

void JFJochImage::SetROIBox(QRect box) {
    roi_type = RoiType::RoiBox;
    roiBox= box;
    roiStartPos = roiBox.topLeft();
    roiEndPos = roiBox.bottomRight();
    Redraw();
}

void JFJochImage::SetROICircle(double x, double y, double radius) {
    roi_type = RoiType::RoiCircle;
    roiBox= QRectF(x - radius, y - radius, 2 * radius, 2 * radius).normalized();
    roiStartPos = roiBox.topLeft();
    roiEndPos = roiBox.bottomRight();
    Redraw();
}

void JFJochImage::mousePressEvent(QMouseEvent *event) {
    if (!scene()) return;

    if (event->button() == Qt::LeftButton) {
        const QPointF scenePos = mapToScene(event->pos());
        active_handle_ = hitTestROIHandle(scenePos, 4.0 / std::sqrt(std::max(1e-4, scale_factor)));

        if (active_handle_ != ResizeHandle::None && active_handle_ != ResizeHandle::Inside) {
            mouse_event_type = MouseEventType::ResizingROI;
            roiStartPos = roiBox.topLeft();
            roiEndPos = roiBox.bottomRight();
            setCursor(Qt::SizeAllCursor);
        } else if (roiBox.contains(scenePos)) {
            mouse_event_type = MouseEventType::MovingROI;
            lastMousePos = event->pos();
            setCursor(Qt::ClosedHandCursor);
        } else if (event->modifiers() & Qt::Modifier::SHIFT) {
            mouse_event_type = MouseEventType::DrawingROI;
            roiStartPos = RoundPoint(scenePos);
            roiEndPos = roiStartPos;
            roi_type = (event->modifiers() & Qt::Modifier::CTRL) ? RoiType::RoiCircle : RoiType::RoiBox;
            setCursor(Qt::CrossCursor);
        } else {
            mouse_event_type = MouseEventType::Panning;
            setCursor(Qt::ClosedHandCursor);
            lastMousePos = event->pos();
        }
    }

    QGraphicsView::mousePressEvent(event);
}
void JFJochImage::mouseMoveEvent(QMouseEvent *event) {
    if (!scene())
        return;

    const QPointF scenePos = mapToScene(event->pos());
    mouseHover(event);
    QPointF delta;

    switch (mouse_event_type) {
        case MouseEventType::Panning:
            delta = mapToScene(event->pos()) - mapToScene(lastMousePos);
            lastMousePos = event->pos();
            translate(delta.x(), delta.y());
            updateOverlay();
            break;
        case MouseEventType::DrawingROI:
            roiEndPos = RoundPoint(scenePos);
            updateROI();
            break;
        case MouseEventType::MovingROI:
            delta = mapToScene(event->pos()) - mapToScene(lastMousePos);
            lastMousePos = event->pos();
            roiBox.translate(delta);
            updateROI();
            break;
        case MouseEventType::ResizingROI: {
            // Modify the corresponding edges based on active_handle_
            if (roi_type == RoiType::RoiCircle) {
                // Resize circle by radius only, keep center fixed
                const QPointF c = roiBox.center();
                const qreal dx = scenePos.x() - c.x();
                const qreal dy = scenePos.y() - c.y();
                qreal r = std::hypot(dx, dy);
                const qreal rMin = 1.0; // clamp tiny radii
                if (r < rMin) r = rMin;
                roiBox = QRectF(c.x() - r, c.y() - r, 2*r, 2*r);
            } else {
                // Box: modify edges based on active handle
                QRectF r = roiBox;
                switch (active_handle_) {
                    case ResizeHandle::Left:       r.setLeft(scenePos.x()); break;
                    case ResizeHandle::Right:      r.setRight(scenePos.x()); break;
                    case ResizeHandle::Top:        r.setTop(scenePos.y()); break;
                    case ResizeHandle::Bottom:     r.setBottom(scenePos.y()); break;
                    case ResizeHandle::TopLeft:    r.setTop(scenePos.y()); r.setLeft(scenePos.x()); break;
                    case ResizeHandle::TopRight:   r.setTop(scenePos.y()); r.setRight(scenePos.x()); break;
                    case ResizeHandle::BottomLeft: r.setBottom(scenePos.y()); r.setLeft(scenePos.x()); break;
                    case ResizeHandle::BottomRight:r.setBottom(scenePos.y()); r.setRight(scenePos.x()); break;
                    default: break;
                }
                roiBox = r.normalized();

            }
            updateROI();
            break;
        }
        case MouseEventType::None: {
            const qreal tol = 4.0 / std::sqrt(std::max(1e-4, scale_factor));
            ResizeHandle h = hitTestROIHandle(scenePos, tol);
            // Update hover state so overlay can draw arrows/handles accordingly
            if (h != hover_handle_) {
                hover_handle_ = h;
                updateOverlay();
            }
            // Set an informative cursor
            switch (h) {
                case ResizeHandle::Left:
                case ResizeHandle::Right:
                    setCursor(Qt::SizeHorCursor); break;
                case ResizeHandle::Top:
                case ResizeHandle::Bottom:
                    setCursor(Qt::SizeVerCursor); break;
                case ResizeHandle::TopLeft:
                case ResizeHandle::BottomRight:
                    setCursor(Qt::SizeFDiagCursor); break;
                case ResizeHandle::TopRight:
                case ResizeHandle::BottomLeft:
                    setCursor(Qt::SizeBDiagCursor); break;
                case ResizeHandle::Inside:
                    setCursor(Qt::OpenHandCursor); break;
                case ResizeHandle::None:
                    setCursor(Qt::ArrowCursor); break;
            }

            break;
        }
    }

    QGraphicsView::mouseMoveEvent(event);
}

void JFJochImage::mouseReleaseEvent(QMouseEvent *event) {
    if (!scene()) return;

    if (event->button() == Qt::LeftButton) {
        if (mouse_event_type == MouseEventType::DrawingROI) {
            roiEndPos = RoundPoint(mapToScene(event->pos()));
        }
        updateROI();
    }

    mouse_event_type = MouseEventType::None;
    active_handle_ = ResizeHandle::None;
    setCursor(Qt::ArrowCursor);
    QGraphicsView::mouseReleaseEvent(event);

}

JFJochImage::ResizeHandle
JFJochImage::hitTestROIHandle(const QPointF& scenePos, qreal tol) const {
        if (roiBox.isNull() || roiBox.width() <= 0 || roiBox.height() <= 0)
        return ResizeHandle::None;

    const QRectF r = roiBox;

    if (roi_type == RoiType::RoiCircle) {
        // Circle hit test: near perimeter -> resize, inside -> move
        const QPointF c = r.center();
        const qreal rx = r.width() * 0.5;
        const qreal ry = r.height() * 0.5;
        // Enforce circular assumption: use average radius
        const qreal rad = 0.5 * (rx + ry);
        const qreal dx = scenePos.x() - c.x();
        const qreal dy = scenePos.y() - c.y();
        const qreal d = std::hypot(dx, dy);
        if (std::abs(d - rad) <= tol) {
            // generic "edge" resize handle for circle
            return ResizeHandle::Right;
        }
        if (d < rad) return ResizeHandle::Inside;
        return ResizeHandle::None;
    }

    // Box hit test (corners first)
    const QPointF tl = r.topLeft();
    const QPointF tr = r.topRight();
    const QPointF bl = r.bottomLeft();
    const QPointF br = r.bottomRight();

    auto nearPt = [&](const QPointF& a, const QPointF& b, qreal t) {
        return std::abs(a.x() - b.x()) <= t && std::abs(a.y() - b.y()) <= t;
    };

    if (nearPt(scenePos, tl, tol)) return ResizeHandle::TopLeft;
    if (nearPt(scenePos, tr, tol)) return ResizeHandle::TopRight;
    if (nearPt(scenePos, bl, tol)) return ResizeHandle::BottomLeft;
    if (nearPt(scenePos, br, tol)) return ResizeHandle::BottomRight;

    // Edges
    if (std::abs(scenePos.x() - r.left()) <= tol && scenePos.y() >= r.top() - tol && scenePos.y() <= r.bottom() + tol)
        return ResizeHandle::Left;
    if (std::abs(scenePos.x() - r.right()) <= tol && scenePos.y() >= r.top() - tol && scenePos.y() <= r.bottom() + tol)
        return ResizeHandle::Right;
    if (std::abs(scenePos.y() - r.top()) <= tol && scenePos.x() >= r.left() - tol && scenePos.x() <= r.right() + tol)
        return ResizeHandle::Top;
    if (std::abs(scenePos.y() - r.bottom()) <= tol && scenePos.x() >= r.left() - tol && scenePos.x() <= r.right() + tol)
        return ResizeHandle::Bottom;

    if (r.contains(scenePos)) return ResizeHandle::Inside;
    return ResizeHandle::None;
}

void JFJochImage::updateROI() {
    if (roi_type == RoiType::RoiBox) {
        if (mouse_event_type == MouseEventType::DrawingROI) {
            // While drawing: construct box from start/end
            QRectF rect = QRectF(RoundPoint(roiStartPos), RoundPoint(roiEndPos)).normalized();
            roiBox = rect;
        } else {
            // While moving/resizing: keep roiBox as modified, just sync corners
            roiStartPos = roiBox.topLeft();
            roiEndPos   = roiBox.bottomRight();
        }
        emit roiBoxUpdated(roiBox.toRect());
    } else {
        if (mouse_event_type == MouseEventType::DrawingROI) {
            // Center at roiStartPos, radius from start->end
            QPointF delta = roiStartPos - roiEndPos;
            double radius = std::sqrt(delta.x() * delta.x() + delta.y() * delta.y());
            roiBox = QRectF(roiStartPos.x() - radius, roiStartPos.y() - radius,
                            2 * radius, 2 * radius).normalized();
            emit roiCircleUpdated(roiStartPos.x(), roiStartPos.y(), radius);
        } else {
            // Moving/resizing: infer center/radius from roiBox
            const QPointF c = roiBox.center();
            const double radius = 0.5 * std::min(roiBox.width(), roiBox.height());
            roiStartPos = c;                          // treat start as center for consistency
            roiEndPos   = QPointF(c.x() + radius, c.y()); // arbitrary point on radius
            emit roiCircleUpdated(c.x(), c.y(), radius);
        }
    }
    updateOverlay();
}

void JFJochImage::DrawROI() {
    if (roiBox.isNull() || roiBox.width() <= 0 || roiBox.height() <= 0) return;

    auto scn = scene();
    if (!scn)
        return;

    QPen pen(feature_color, 2);
    pen.setStyle(Qt::DashLine);
    pen.setCosmetic(true);

    const qreal f = std::clamp(scale_factor, 0.5, 50.0);
    const qreal handleSize = 3.0 / std::sqrt(std::max(1e-4, f));

    if (roi_type == RoiType::RoiCircle) {
        // Draw circle
        scn->addEllipse(roiBox, pen);

        // A single handle on the circle at the rightmost point
        const QPointF c = roiBox.center();
        const qreal rad = 0.5 * (roiBox.width() + roiBox.height()) * 0.5; // average, should be equal
        QPointF hpos = QPointF(roiBox.right(), c.y());
        scn->addRect(QRectF(hpos.x() - handleSize, hpos.y() - handleSize, 2 * handleSize, 2 * handleSize),
                     QPen(feature_color, 1), QBrush(feature_color));

        // On hover near perimeter: draw in/out arrows along radius at handle
        if (hover_handle_ != ResizeHandle::None && hover_handle_ != ResizeHandle::Inside) {
            QPen apen(feature_color, 1);
            apen.setCosmetic(true);
            const qreal arrowLen = 8.0 / std::sqrt(std::max(1e-4, f));
            // Outward arrow
            scn->addLine(QLineF(c, c + QPointF(rad + arrowLen, 0)), apen);
            // Inward arrow
            scn->addLine(QLineF(c, c + QPointF(rad - arrowLen, 0)), apen);
        }
    } else {
        // Box
        scn->addRect(roiBox, pen);

        // Corner handles
        auto addHandle = [&](const QPointF& p) {
            scn->addRect(QRectF(p.x() - handleSize, p.y() - handleSize, 2 * handleSize, 2 * handleSize),
                         QPen(feature_color, 1), QBrush(feature_color));
        };
        addHandle(roiBox.topLeft());
        addHandle(roiBox.topRight());
        addHandle(roiBox.bottomLeft());
        addHandle(roiBox.bottomRight());

        // On hover over a resizable edge/corner: draw small arrows indicating resize direction
        if (hover_handle_ != ResizeHandle::None && hover_handle_ != ResizeHandle::Inside) {
            QPen apen(feature_color, 1);
            apen.setCosmetic(true);
            const qreal arrowLen = 6.0 / std::sqrt(std::max(1e-4, f));
            const qreal off = 10.0 / std::sqrt(std::max(1e-4, f));
            auto drawArrow = [&](const QPointF& a, const QPointF& b) {
                scn->addLine(QLineF(a, b), apen);
            };
            const QRectF r = roiBox;
            switch (hover_handle_) {
                case ResizeHandle::Left:
                    drawArrow(QPointF(r.left(), r.center().y() - off), QPointF(r.left() - arrowLen, r.center().y() - off));
                    drawArrow(QPointF(r.left(), r.center().y() + off), QPointF(r.left() - arrowLen, r.center().y() + off));
                    break;
                case ResizeHandle::Right:
                    drawArrow(QPointF(r.right(), r.center().y() - off), QPointF(r.right() + arrowLen, r.center().y() - off));
                    drawArrow(QPointF(r.right(), r.center().y() + off), QPointF(r.right() + arrowLen, r.center().y() + off));
                    break;
                case ResizeHandle::Top:
                    drawArrow(QPointF(r.center().x() - off, r.top()), QPointF(r.center().x() - off, r.top() - arrowLen));
                    drawArrow(QPointF(r.center().x() + off, r.top()), QPointF(r.center().x() + off, r.top() - arrowLen));
                    break;
                case ResizeHandle::Bottom:
                    drawArrow(QPointF(r.center().x() - off, r.bottom()), QPointF(r.center().x() - off, r.bottom() + arrowLen));
                    drawArrow(QPointF(r.center().x() + off, r.bottom()), QPointF(r.center().x() + off, r.bottom() + arrowLen));
                    break;
                case ResizeHandle::TopLeft:
                case ResizeHandle::TopRight:
                case ResizeHandle::BottomLeft:
                case ResizeHandle::BottomRight:
                    // For corners, show arrows on both axes (simple version)
                    drawArrow(QPointF(r.right(), r.center().y()), QPointF(r.right() + arrowLen, r.center().y()));
                    drawArrow(QPointF(r.left(),  r.center().y()), QPointF(r.left() - arrowLen,  r.center().y()));
                    drawArrow(QPointF(r.center().x(), r.top()), QPointF(r.center().x(), r.top() - arrowLen));
                    drawArrow(QPointF(r.center().x(), r.bottom()), QPointF(r.center().x(), r.bottom() + arrowLen));
                    break;
                default: break;
            }
        }
    }
}

void JFJochImage::Redraw() {
    if (W*H <= 0)
        return;

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

    QGraphicsScene *currentScene = scene();
    if (!currentScene) {
        // First time - create a new scene
        currentScene = new QGraphicsScene(this);
        setScene(currentScene);
    }

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

void JFJochImage::GeneratePixmap() {

    /*
    if (auto_fg || auto_bg) {
        auto vec1(image_fp);
        std::sort(vec1.begin(), vec1.end());

        if (auto_fg) {
            foreground = vec1[vec1.size() - 1];
            emit foregroundChanged(foreground);
        }
        if (auto_bg) {
            background = vec1[0];
            emit backgroundChanged(background);
        }
    } */

    QImage qimg(int(W), int(H), QImage::Format_RGB888);
    image_rgb.resize(W * H);

    // Bad pixel 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)};

    // Saturation color
    rgb sat_color{};
    if (show_saturation) {
        sat_color = bad_color;
    } else
        sat_color = color_scale.Apply(1.0);

    for (int y = 0; y < H; ++y) {
        uchar *scanLine = qimg.scanLine(y); // Get writable pointer to the row
        for (int x = 0; x < W; ++x) {
            float fp = image_fp[y * W + x];
            if (std::isnan(fp))
                image_rgb[y * W + x] = color_scale.Apply(ColorScaleSpecial::Gap);
            else if (std::isinf(fp)) {
                if (std::signbit(fp))
                    image_rgb[y * W + x] = bad_color;
                else
                    image_rgb[y * W + x] = sat_color;
            } else
                image_rgb[y * W + x] = color_scale.Apply(image_fp[y * W + x], background, foreground);
            scanLine[x * 3 + 0] = image_rgb[y * W + x].r; // Red
            scanLine[x * 3 + 1] = image_rgb[y * W + x].g; // Green
            scanLine[x * 3 + 2] = image_rgb[y * W + x].b; // Blue
        }
    }

    pixmap = QPixmap::fromImage(qimg);
}

void JFJochImage::centerOnSpot(QPointF point) {
    // If W or H = 0, then conditions are never satisfied
    if (point.x() >= 0 && point.x() < W && point.y() >= 0 && point.y() < H)
        centerOn(point);
}

void JFJochImage::writePixelLabels() {

    static QFont font([] {
        QFont f("DejaVu Sans Mono");
        f.setStyleHint(QFont::TypeWriter);
        f.setPixelSize(1);
        return f;
    }());
    static const QString kGap = QStringLiteral("Gap");
    static const QString kErr = QStringLiteral("Err");
    static const QString kSat = QStringLiteral("Sat");

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

    const int startX = std::max(0, static_cast<int>(std::floor(visibleRect.left())));
    const int endX   = std::min(static_cast<int>(W), static_cast<int>(std::ceil(visibleRect.right())));
    const int startY = std::max(0, static_cast<int>(std::floor(visibleRect.top())));
    const int endY   = std::min(static_cast<int>(H), static_cast<int>(std::ceil(visibleRect.bottom())));

    const int visW = std::max(0, endX - startX);
    const int visH = std::max(0, endY - startY);
    int maxLabels = 1000;

    if (visW * visH  <= maxLabels) {
        QString numBuf; // reused buffer

        for (int y = startY; y < endY; y ++) {
            for (int x = startX; x < endX; x++) {
                const int idx = y * W + x;
                const float val = image_fp[idx];

                const QString* pText = nullptr;
                if (std::isnan(val)) {
                    pText = &kGap;
                } else if (std::isinf(val)) {
                    pText = std::signbit(val) ? &kErr : &kSat;
                } else {
                    // Fixed format reduces overhead and string length variability
                    numBuf = QString::number(val, 'g', 4);
                    pText = &numBuf;
                }

                QGraphicsTextItem* textItem = scene()->addText(*pText, font);
                if (luminance(image_rgb[idx]) > 128.0)
                    textItem->setDefaultTextColor(Qt::black);
                else
                    textItem->setDefaultTextColor(Qt::white);

                textItem->setPos(x - 0.7, y - 0.8);
                textItem->setScale(0.2);
            }
        }
    }
}

void JFJochImage::updateOverlay() {
    if (!scene() || W*H <= 0) return;

    scene()->clear();
    scene()->addItem(new QGraphicsPixmapItem(pixmap));

    if (scale_factor > 30.0)
        writePixelLabels();


    DrawROI();

    addCustomOverlay();
}

void JFJochImage::addCustomOverlay() {}