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

#include "JFJochViewerImageStrip.h"

#include <QComboBox>
#include <QCheckBox>
#include <QHBoxLayout>
#include <QVBoxLayout>
#include <QScrollArea>
#include <QToolButton>
#include <QLabel>
#include <QPixmap>

#include <algorithm>
#include <numeric>
#include <cmath>

JFJochViewerImageStrip::JFJochViewerImageStrip(QWidget *parent) : QWidget(parent) {
    auto *layout = new QVBoxLayout(this);
    layout->setContentsMargins(4, 2, 4, 2);

    auto *controls = new QHBoxLayout();
    controls->addWidget(new QLabel("Show", this));
    mode_ = new QComboBox(this);
    mode_->addItem("Evenly spaced", 0);
    mode_->addItem("Most spots", 1);
    mode_->addItem("Indexed", 2);
    mode_->addItem("Resolution", 3);     // representatives across the resolution range
    mode_->addItem("Background", 4);      // representatives across the background range
    controls->addWidget(mode_);
    spots_ = new QCheckBox("Spots", this);
    spots_->setChecked(true);
    spots_->setToolTip("Overlay found spots — a real pattern reads like a constellation");
    controls->addWidget(spots_);
    controls->addStretch();
    layout->addLayout(controls);

    auto *scroll = new QScrollArea(this);
    scroll->setWidgetResizable(true);
    scroll->setVerticalScrollBarPolicy(Qt::ScrollBarAlwaysOff);
    scroll->setHorizontalScrollBarPolicy(Qt::ScrollBarAsNeeded);
    auto *stripWidget = new QWidget(scroll);
    strip_ = new QHBoxLayout(stripWidget);
    strip_->setContentsMargins(0, 0, 0, 0);
    strip_->addStretch();
    scroll->setWidget(stripWidget);
    layout->addWidget(scroll);

    connect(mode_, &QComboBox::currentIndexChanged, this, [this] { Rebuild(); });
    connect(spots_, &QCheckBox::toggled, this, [this] { Rebuild(); });
}

void JFJochViewerImageStrip::datasetLoaded(std::shared_ptr<const JFJochReaderDataset> dataset) {
    dataset_ = std::move(dataset);
    Rebuild();
}

void JFJochViewerImageStrip::thumbnailReady(qint64 image_number, QImage thumb) {
    auto it = slots_.find(image_number);
    if (it != slots_.end())
        it.value()->setIcon(QPixmap::fromImage(thumb));
}

QVector<qint64> JFJochViewerImageStrip::ComputeRepresentatives() const {
    QVector<qint64> result;
    if (!dataset_)
        return result;
    const int64_t total = dataset_->experiment.GetImageNum();
    if (total <= 0)
        return result;
    const int N = static_cast<int>(std::min<int64_t>(8, total));
    const int mode = mode_->currentIndex();
    const auto &sc = dataset_->spot_count;
    const auto &sci = dataset_->spot_count_indexed;

    auto evenly = [&] {
        for (int i = 0; i < N; ++i)
            result.push_back(N == 1 ? 0 : static_cast<int64_t>(i) * (total - 1) / (N - 1));
    };

    // Pick N images that span a metric's distribution (sort by value, sample evenly across it) —
    // a quick "histogram representative" selection.
    auto spread = [&](const std::vector<float> &metric) {
        std::vector<std::pair<float, int64_t>> vals;
        for (size_t i = 0; i < metric.size(); ++i)
            if (std::isfinite(metric[i]))
                vals.push_back({metric[i], static_cast<int64_t>(i)});
        if (vals.empty()) { evenly(); return; }
        std::sort(vals.begin(), vals.end());
        const int n = static_cast<int>(std::min<size_t>(N, vals.size()));
        for (int i = 0; i < n; ++i)
            result.push_back(vals[n == 1 ? 0 : static_cast<size_t>(i) * (vals.size() - 1) / (n - 1)].second);
    };

    if (mode == 1 && !sc.empty()) {                 // most spots
        std::vector<int64_t> idx(sc.size());
        std::iota(idx.begin(), idx.end(), 0);
        const size_t n = std::min<size_t>(N, idx.size());
        std::partial_sort(idx.begin(), idx.begin() + n, idx.end(),
                          [&](int64_t a, int64_t b) { return sc[a] > sc[b]; });
        idx.resize(n);
        std::sort(idx.begin(), idx.end());
        for (int64_t i : idx) result.push_back(i);
    } else if (mode == 2 && !sci.empty()) {         // indexed images, evenly sampled
        std::vector<int64_t> indexed;
        for (size_t i = 0; i < sci.size(); ++i)
            if (sci[i] > 0) indexed.push_back(static_cast<int64_t>(i));
        if (indexed.empty()) {
            evenly();
        } else {
            const int n = static_cast<int>(std::min<size_t>(N, indexed.size()));
            for (int i = 0; i < n; ++i)
                result.push_back(indexed[n == 1 ? 0 : static_cast<size_t>(i) * (indexed.size() - 1) / (n - 1)]);
        }
    } else if (mode == 3 && !dataset_->resolution_estimate.empty()) {   // resolution spread
        spread(dataset_->resolution_estimate);
    } else if (mode == 4 && !dataset_->bkg_estimate.empty()) {          // background spread
        spread(dataset_->bkg_estimate);
    } else {
        evenly();
    }
    return result;
}

void JFJochViewerImageStrip::Rebuild() {
    // Clear the strip (keep the trailing stretch).
    slots_.clear();
    while (strip_->count() > 1) {
        QLayoutItem *item = strip_->takeAt(0);
        if (item->widget())
            item->widget()->deleteLater();
        delete item;
    }
    if (!dataset_)
        return;

    const QVector<qint64> reps = ComputeRepresentatives();
    int pos = 0;
    for (qint64 n : reps) {
        auto *btn = new QToolButton(this);
        btn->setToolButtonStyle(Qt::ToolButtonTextUnderIcon);
        btn->setIconSize(QSize(120, 120));
        btn->setText(QString::number(n + 1));
        btn->setAutoRaise(true);
        btn->setCursor(Qt::PointingHandCursor);
        btn->setToolTip(QStringLiteral("Open image %1").arg(n + 1));
        connect(btn, &QToolButton::clicked, this, [this, n] { emit imageSelected(n, 1); });
        strip_->insertWidget(pos++, btn);
        slots_[n] = btn;
    }
    emit requestThumbnails(reps, spots_->isChecked());
}