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

#include "JFJochSimpleImage.h"
#include <QGraphicsScene>
#include <QGraphicsPixmapItem>
#include <QWheelEvent>
#include <QScrollBar>
#include <QtMath>
#include <cstring>
#include <algorithm>

#include "../../common/JFJochException.h"

static inline double normalize_to_unit(double raw, double bg, double fg) {
    const double denom = std::max(1e-12, double(fg) - double(bg));
    double v = (raw - bg) / denom;
    if (v < 0.0) v = 0.0;
    if (v > 1.0) v = 1.0;
    return v;
}

JFJochSimpleImage::JFJochSimpleImage(QWidget *parent)
    : JFJochImage(parent) {
    auto *scn = new QGraphicsScene(this);
    setScene(scn);

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

    // Keep overlays in pixel units independent of zoom (for labels font sizing)
    setViewportUpdateMode(QGraphicsView::FullViewportUpdate);
}

void JFJochSimpleImage::clear() {
    has_image_ = false;
    image_.reset();

    if (scene())
        scene()->clear();
}

void JFJochSimpleImage::setImage(std::shared_ptr<const SimpleImage> img) {
    image_ = std::move(img);
    has_image_ = true;
    renderImage();
    updateOverlay();
}

void JFJochSimpleImage::Redraw() {
    if (has_image_) {
        renderImage();
        // Preserve current transform while updating
        updateOverlay();
    }
}

void JFJochSimpleImage::mousePressEvent(QMouseEvent* event) {
    if (!scene() || !has_image_) return;

    if (event->button() == Qt::LeftButton) {
        const QPointF scenePos = mapToScene(event->pos());

        // With Shift -> ROI interactions (draw/resize/move); without -> pan
        if (event->modifiers() & Qt::Modifier::SHIFT) {
            // Determine if we are over an existing ROI handle
            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;
                roi_start_pos_ = roi_box_.topLeft();
                roi_end_pos_   = roi_box_.bottomRight();
                setCursor(Qt::SizeAllCursor);
            } else if (roi_box_.contains(scenePos)) {
                mouse_event_type_ = MouseEventType::MovingROI;
                last_mouse_pos_ = event->pos();
                setCursor(Qt::ClosedHandCursor);
            } else {
                mouse_event_type_ = MouseEventType::DrawingROI;
                roi_start_pos_ = RoundPoint(scenePos);
                roi_end_pos_   = roi_start_pos_;
                setCursor(Qt::CrossCursor);
            }
        } else {
            mouse_event_type_ = MouseEventType::Panning;
            panning_ = true;
            setCursor(Qt::ClosedHandCursor);
            last_mouse_pos_ = event->pos();
        }
    }

    QGraphicsView::mousePressEvent(event);
}

void JFJochSimpleImage::mouseMoveEvent(QMouseEvent* event) {
    if (!scene() || !has_image_) return;

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

    switch (mouse_event_type_) {
        case MouseEventType::Panning: {
            const QPointF delta = mapToScene(event->pos()) - mapToScene(last_mouse_pos_);
            last_mouse_pos_ = event->pos();
            translate(delta.x(), delta.y());
            updateOverlay();
            break;
        }
        case MouseEventType::DrawingROI: {
            roi_end_pos_ = RoundPoint(scenePos);
            updateROI();
            break;
        }
        case MouseEventType::MovingROI: {
            const QPointF delta = mapToScene(event->pos()) - mapToScene(last_mouse_pos_);
            last_mouse_pos_ = event->pos();
            roi_box_.translate(delta);
            updateROI();
            break;
        }
        case MouseEventType::ResizingROI: {
            // Modify the corresponding edges based on active_handle_
            QRectF r = roi_box_;
            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;
            }
            roi_box_ = r.normalized();
            updateROI();
            break;
        }
        case MouseEventType::None: {
            // Hover feedback / status bar display
            if ((scenePos.x() >= 0)
                && (scenePos.x() < image_->image.GetWidth())
                && (scenePos.y() >= 0)
                && (scenePos.y() < image_->image.GetHeight())) {
                const auto ix = int(scenePos.x());
                const auto iy = int(scenePos.y());
                const auto idx = iy * int(image_->image.GetWidth()) + ix;
                if (idx >= 0 && idx < int(image_values_.size()))
                    emit writeStatusBar(QString("x=%1 y=%2 I=%3")
                                        .arg(scenePos.x(), 0, 'f', 1)
                                        .arg(scenePos.y(), 0, 'f', 1)
                                        .arg(image_values_[size_t(idx)]), 3000);
            } else {
                emit writeStatusBar("", 1000);
            }
            break;
        }
    }

    QGraphicsView::mouseMoveEvent(event);
}

void JFJochSimpleImage::mouseReleaseEvent(QMouseEvent* event) {
    if (!scene() || !has_image_) return;

    if (event->button() == Qt::LeftButton) {
        if (mouse_event_type_ == MouseEventType::DrawingROI) {
            roi_end_pos_ = RoundPoint(mapToScene(event->pos()));
            updateROI();
        }
        panning_ = false;
        mouse_event_type_ = MouseEventType::None;
        active_handle_ = ResizeHandle::None;
        setCursor(Qt::ArrowCursor);
    }

    QGraphicsView::mouseReleaseEvent(event);
}


void JFJochSimpleImage::resizeEvent(QResizeEvent *event) {
    QGraphicsView::resizeEvent(event);
    // Nothing special; keep current transform and scene
}

void JFJochSimpleImage::renderImage() {
    const size_t W = image_->image.GetWidth();
    const size_t H = image_->image.GetHeight();
    if (W == 0 || H == 0) return;

    // Prepare destination QImage (RGB888)
    QImage qimg(int(W), int(H), QImage::Format_RGB888);
    image_rgb.resize(W * H);
    image_values_.resize(W * H);

    std::vector<uint8_t> image_buffer;

    // Access uncompressed data
    const uint8_t *src = image_->image.GetUncompressedPtr(image_buffer);
    const auto mode = image_->image.GetMode();

    auto write_rgb_row = [&](int y) -> uchar * {
        return qimg.scanLine(y);
    };

    auto write_pixel = [&](size_t idx, uchar *dstRow, const rgb &c) {
        dstRow[idx * 3 + 0] = c.r;
        dstRow[idx * 3 + 1] = c.g;
        dstRow[idx * 3 + 2] = c.b;
        image_rgb[idx] = c;
    };

    // Fast path: packed RGB (assumed 3 bytes per pixel)
    if (mode == CompressedImageMode::RGB) {
        for (int y = 0; y < int(H); ++y) {
            const rgb *row = reinterpret_cast<const rgb *>(src + y * W * sizeof(rgb));
            uchar *dst = write_rgb_row(y);
            for (size_t x = 0; x < W; ++x) {
                const rgb c = row[x];
                write_pixel(x, dst, c);
            }
        }
    } else {
        switch (mode) {
            case CompressedImageMode::Uint8:
                renderImage<uint8_t>(qimg, src);
                break;
            case CompressedImageMode::Int8:
                renderImage<int8_t>(qimg, src);
                break;
            case CompressedImageMode::Uint16:
                renderImage<uint16_t>(qimg, src);
                break;
            case CompressedImageMode::Int16:
                renderImage<int16_t>(qimg, src);
                break;
            case CompressedImageMode::Uint32:
                renderImage<uint32_t>(qimg, src);
                break;
            case CompressedImageMode::Int32:
                renderImage<int32_t>(qimg, src);
                break;
            case CompressedImageMode::Float32:
                renderImage<float>(qimg, src);
                break;
            case CompressedImageMode::Float64:
                renderImage<double>(qimg, src);
                break;
            case CompressedImageMode::Float16:
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Float16 not supported");
                break;
        }
    }

    // Build pixmap
    pixmap = QPixmap::fromImage(qimg);
}

void JFJochSimpleImage::updateOverlay() {
    if (!scene()) return;
    scene()->clear();
    if (!has_image_) return;

    // Add base image
    scene()->addItem(new QGraphicsPixmapItem(pixmap));

    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
    const int W = int(image_->image.GetWidth());
    const int H = int(image_->image.GetHeight());

    const int startX = std::max(0, static_cast<int>(std::floor(visibleRect.left())));
    const int endX = std::min(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(H, static_cast<int>(std::ceil(visibleRect.bottom())));

    // Only when zoomed in and region small (performance)
    if (scale_factor > 20.0 && (endX - startX + 1) * (endY - startY + 1) < 500) {
        const qreal f = std::clamp(scale_factor, 0.5, 50.0);
        QFont font("Arial", 1);
        font.setPixelSize(1); // 1 px font; we use scaling below as needed

        for (int y = startY; y < endY; ++y) {
            const size_t base = size_t(y) * size_t(W);
            for (int x = startX; x < endX; ++x) {
                const size_t idx = base + size_t(x);

                const double raw = image_values_[idx];
                const rgb c = image_rgb[idx];

                // If no scalar was stored (e.g., RGB source), you can skip or compute some alt text
                if (std::isnan(raw)) continue;

                // Text content from original value
                const QString pixelText = QString::number(raw);

                // Contrast against colored pixel
                QGraphicsTextItem *textItem = scene()->addText(pixelText, font);
                if (luminance(c) > 128.0)
                    textItem->setDefaultTextColor(Qt::black);
                else
                    textItem->setDefaultTextColor(Qt::white);

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

    drawROI(scene());
}

template<class T>
void JFJochSimpleImage::renderImage(QImage &qimg, const uint8_t *input) {
    const size_t W = image_->image.GetWidth();
    const size_t H = image_->image.GetHeight();

    auto ptr = reinterpret_cast<const T *>(input);
    int64_t val;
    for (int i = 0; i < W * H; i++) {
        if (std::is_floating_point_v<T>)
            val = std::lround(ptr[i]);
        else
            val = static_cast<int64_t>(ptr[i]);
        image_values_[i] = val;
    }

    if (auto_fgbg) {
        auto vec1(image_values_);
        std::sort(vec1.begin(), vec1.end());
        background = vec1[0];
        foreground = vec1[vec1.size() - 1];
        emit backgroundChanged(background);
        emit foregroundChanged(foreground);
    }

    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) {
            image_rgb[y * W + x] = color_scale.Apply(image_values_[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
        }
    }
}


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

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

    const QRectF r = roi_box_;
    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 JFJochSimpleImage::updateROI() {
    // Normalize box if drawing
    if (mouse_event_type_ == MouseEventType::DrawingROI) {
        roi_box_ = QRectF(roi_start_pos_, roi_end_pos_).normalized();
    }

    // Clamp ROI to image bounds
    if (has_image_) {
        const QRectF bounds(0.0, 0.0,
                            image_->image.GetWidth(),
                            image_->image.GetHeight());
        roi_box_ = roi_box_.intersected(bounds);
    }

   scheduleSceneUpdate();
}

void JFJochSimpleImage::drawROI(QGraphicsScene* scn) {
    if (roi_box_.isNull() || roi_box_.width() <= 0 || roi_box_.height() <= 0) return;

    QPen pen(feature_color, 2);
    pen.setStyle(Qt::DashLine);
    pen.setCosmetic(true);
    scn->addRect(roi_box_, pen);

    // Small corner handles
    const qreal handleSize = 3.0 / std::sqrt(std::max(1e-4, scale_factor));
    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(roi_box_.topLeft());
    addHandle(roi_box_.topRight());
    addHandle(roi_box_.bottomLeft());
    addHandle(roi_box_.bottomRight());

    // Compute stats inside integer pixel ROI
    QRect roi_px = roi_box_.toRect().normalized();
    ROIStats stats = computeROIStats(roi_px);

    if (stats.valid()) {
        // Compose text
        QString text = QString("min %1  max %2  avg %3  std %4  n=%5")
                .arg(stats.min, 0, 'f', 2)
                .arg(stats.max, 0, 'f', 2)
                .arg(stats.avg, 0, 'f', 2)
                .arg(stats.stddev, 0, 'f', 2)
                .arg(qulonglong(stats.count));

        // Choose background based on local luminance under top-left pixel of ROI
        const int W = int(image_->image.GetWidth());
        const int H = int(image_->image.GetHeight());
        int px = std::clamp(roi_px.left(), 0, std::max(0, W - 1));
        int py = std::clamp(roi_px.top(), 0, std::max(0, H - 1));
        const rgb bgc = image_rgb[size_t(py) * size_t(W) + size_t(px)];
        const bool darkText = luminance(bgc) > 128.0;

        QFont font("Arial");
        const qreal f = std::clamp(scale_factor, 0.5, 50.0);
        font.setPointSizeF(14.0 / std::sqrt(f));

        auto* t = scn->addText(text, font);
        t->setDefaultTextColor(darkText ? Qt::black : Qt::white);
        t->setZValue(20.0);

        // Place near top-left of ROI, slightly above
        QPointF pos = roi_box_.topLeft() + QPointF(2.0, -16.0 / std::sqrt(f));
        t->setPos(pos);
    }
}

JFJochSimpleImage::ROIStats
JFJochSimpleImage::computeROIStats(const QRect& roi_px) const {
    ROIStats s;
    if (!has_image_) return s;

    const int W = int(image_->image.GetWidth());
    const int H = int(image_->image.GetHeight());

    const int x0 = std::clamp(roi_px.left(), 0, std::max(0, W - 1));
    const int x1 = std::clamp(roi_px.right(), 0, std::max(0, W - 1));
    const int y0 = std::clamp(roi_px.top(), 0, std::max(0, H - 1));
    const int y1 = std::clamp(roi_px.bottom(), 0, std::max(0, H - 1));
    if (x1 < x0 || y1 < y0) return s;

    // One pass: min, max, sum, sum of squares
    double sum = 0.0;
    double sumsq = 0.0;
    s.min = std::numeric_limits<double>::infinity();
    s.max = -std::numeric_limits<double>::infinity();
    s.count = 0;

    for (int y = y0; y <= y1; ++y) {
        const size_t base = size_t(y) * size_t(W);
        for (int x = x0; x <= x1; ++x) {
            const size_t idx = base + size_t(x);
            const double v = image_values_[idx];
            if (std::isnan(v)) continue; // skip invalid
            if (v < s.min) s.min = v;
            if (v > s.max) s.max = v;
            sum += v;
            sumsq += v * v;
            ++s.count;
        }
    }

    if (!s.valid()) return s;

    s.avg = sum / double(s.count);
    double variance = sumsq / double(s.count) - s.avg * s.avg;
    if (variance < 0.0) variance = 0.0; // guard small negatives from FP error
    s.stddev = std::sqrt(variance);
    return s;
}


void JFJochSimpleImage::onRepaintTimer() {
    // Apply any pending pan
    if (!qFuzzyIsNull(pan_accum_.x()) || !qFuzzyIsNull(pan_accum_.y())) {
        translate(pan_accum_.x(), pan_accum_.y());
        pan_accum_ = QPointF(0.0, 0.0);
    }

    // Rebuild scene if requested
    if (needs_scene_update_) {
        updateOverlay();
        needs_scene_update_ = false;
    }

    // Stop timer if nothing pending
    if (qFuzzyIsNull(pan_accum_.x()) && qFuzzyIsNull(pan_accum_.y()) && !needs_scene_update_) {
        repaint_timer_.stop();
    }
}

void JFJochSimpleImage::schedulePanDelta(const QPointF& d) {
    pan_accum_ += d;
    if (!repaint_timer_.isActive())
        repaint_timer_.start();
}

void JFJochSimpleImage::scheduleSceneUpdate() {
    needs_scene_update_ = true;
    if (!repaint_timer_.isActive())
        repaint_timer_.start();
}