Jungfraujoch/viewer/JFJochViewerDatasetInfo.cpp

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

#include <QGridLayout>
#include <QPushButton>
#include <QColor>

#include "JFJochViewerDatasetInfo.h"

namespace {
    // Stable colour per run by its position in the run list (Original is index 0 -> blue), so a
    // run keeps its colour whether it is the active (primary) line or an overlay.
    QColor RunColor(int index) {
        static const QColor palette[] = {
            QColor(0x1f, 0x77, 0xb4), QColor(0xff, 0x7f, 0x0e), QColor(0x2c, 0xa0, 0x2c),
            QColor(0xd6, 0x27, 0x28), QColor(0x94, 0x67, 0xbd), QColor(0x8c, 0x56, 0x4b),
            QColor(0xe3, 0x77, 0xc2), QColor(0x7f, 0x7f, 0x7f),
        };
        constexpr int n = static_cast<int>(sizeof(palette) / sizeof(palette[0]));
        return palette[((index % n) + n) % n];
    }

    // Combo value of the composite grid-scan map. Well clear of the running metric numbers and
    // below the ROI block (100 and up), so a metric added later cannot collide with it.
    constexpr int kCompositeMetric = 50;

    // A run's image->original-number map as floats for the chart x-axis (empty => identity).
    std::vector<float> XForRun(const JFJochReaderDataset &ds) {
        return {ds.source_image_number.begin(), ds.source_image_number.end()};
    }
}

JFJochViewerDatasetInfo::JFJochViewerDatasetInfo(QWidget *parent) : QWidget(parent) {
    auto layout = new QGridLayout(this);
    combo_box = new QComboBox(this);

    last_selection = 0;

    layout->addWidget(combo_box, 0, 0);

    auto reset_button = new QPushButton("Reset zoom", this);
    reset_button->setFixedWidth(100);

    grid_button = new QPushButton("Grid", this);
    grid_button->setFixedWidth(100);
    grid_button->setCheckable(true);
    grid_button->setEnabled(false);

    image_button = new QPushButton("Per-image", this);
    image_button->setFixedWidth(100);
    image_button->setCheckable(true);

    auto add_plot_button = new QPushButton("+ Plot", this);
    add_plot_button->setFixedWidth(80);
    add_plot_button->setToolTip("Open another dataset-info plot, side by side");
    add_plot_button->setStyleSheet(
        "QPushButton { background-color:#1F3A5F; color:white; font-weight:bold; border:none;"
        " border-radius:3px; padding:3px 8px; } QPushButton:hover { background-color:#16314F; }");
    connect(add_plot_button, &QPushButton::clicked, this, &JFJochViewerDatasetInfo::addPlot);

    layout->addWidget(reset_button, 0, 2);
    layout->addWidget(grid_button, 0, 3);
    layout->addWidget(image_button, 0, 4);
    layout->addWidget(add_plot_button, 0, 5);

    stack = new QStackedWidget(this);
    chart_view = new JFJochDatasetInfoChartView(this);
    chart_view->setMinimumHeight(140);   // keep enough room for the axis labels even when compressed
    grid_scan_image = new JFJochGridScanImage(this);
    image_chart = new JFJochViewerSidePanelChart(this);   // per-image profiles, folded in from the side panel

    stack->addWidget(chart_view);
    stack->addWidget(grid_scan_image);
    stack->addWidget(image_chart);

    layout->addWidget(stack, 1, 0, 1, 6);

    connect(chart_view, &JFJochDatasetInfoChartView::imageSelected,
            this, &JFJochViewerDatasetInfo::imageSelectedInChart);

    connect(reset_button, &QPushButton::clicked,
            this, &JFJochViewerDatasetInfo::resetZoomButtonPressed);

    connect(combo_box, &QComboBox::currentIndexChanged,
            this, &JFJochViewerDatasetInfo::comboBoxSelected);

    connect(grid_scan_image, &JFJochGridScanImage::imageSelected,
            this, &JFJochViewerDatasetInfo::imageSelectedInChart);

    connect(chart_view, &JFJochDatasetInfoChartView::writeStatusBar,
            [this](QString string, int timeout_ms) {
                emit writeStatusBar(string, timeout_ms);
            });

    connect(grid_scan_image, &JFJochGridScanImage::writeStatusBar,
            [this](QString string, int timeout_ms) {
                emit writeStatusBar(string, timeout_ms);
            });

    connect(grid_button, &QPushButton::clicked,
            [this]() {
                if (grid_button->isChecked()) {
                    image_button->setChecked(false);
                    combo_box->setEnabled(true);
                }
                UpdateView();
            });
    connect(image_button, &QPushButton::clicked,
            [this]() {
                if (image_button->isChecked())
                    grid_button->setChecked(false);
                UpdateLabels();                  // swap the one combo's contents for the mode
                combo_box->setCurrentIndex(0);
                comboBoxSelected(0);             // apply the first option of the new mode
                UpdateView();
            });
    connect(image_chart, &JFJochViewerSidePanelChart::writeStatusBar,
            [this](QString s, int t) { emit writeStatusBar(s, t); });
    setLayout(layout);
}

void JFJochViewerDatasetInfo::RestoreSelection() {
    QSignalBlocker bl(combo_box);
    combo_box->setCurrentIndex(last_selection < combo_box->count() ? last_selection : 0);
}

void JFJochViewerDatasetInfo::UpdateLabels() {
    QSignalBlocker bl(combo_box);

    if (combo_box->count() > 0)
        last_selection = combo_box->currentIndex();

    combo_box->clear();

    // In per-image mode the one combo selects which per-image profile to show (no second combo).
    if (image_button && image_button->isChecked()) {
        for (const auto &pt : JFJochViewerSidePanelChart::PlotTypes())
            combo_box->addItem(pt.first, pt.second);
        return;
    }

    if (this->dataset) {
        // First for a grid scan: the composite is what a raster is looked at for, and the whole
        // point is that it comes up without anyone having to know to ask for it.
        if (dataset->experiment.GetGridScan())
            combo_box->addItem("Protein / ice map", kCompositeMetric);
        combo_box->addItem("Background estimate", 0);
        combo_box->addItem("Ice ring ratio", 14);
        combo_box->addItem("Spindle blind fraction", 15);
        combo_box->addItem("Resolution estimate", 7);
        combo_box->addItem("Spot count", 1);
        combo_box->addItem("Spot count (indexed)", 2);
        combo_box->addItem("Spot count (ice rings)", 3);
        combo_box->addItem("Spot count (low res.)", 4);

        // Offer indexing/scaling metrics if the active run, any overlay run, or the live run has them.
        bool has_indexing = !dataset->indexing_result.empty();
        bool has_scale = !dataset->image_scale_factor.empty();
        for (const auto &r: runs_)
            if (r.dataset) {
                has_indexing = has_indexing || !r.dataset->indexing_result.empty();
                has_scale = has_scale || !r.dataset->image_scale_factor.empty();
            }
        if (live_run_) {
            has_indexing = has_indexing || !live_run_->indexing_result.empty();
            has_scale = has_scale || !live_run_->image_scale_factor.empty();
        }

        if (has_indexing) {
            combo_box->insertSeparator(1000);
            combo_box->addItem("Indexing result", 5);
            combo_box->addItem("Profile radius", 6);
            combo_box->addItem("B-factor", 8);
            combo_box->addItem("Mosaicity", 9);
            combo_box->addItem("Integrated reflections", 12);
            combo_box->addItem("Indexing lattice count", 13);
        }

        if (has_scale) {
            combo_box->insertSeparator(1000);
            combo_box->addItem("Scale factor", 10);
            combo_box->addItem("CC", 11);
        }
        for (int i = 0; i < this->dataset->roi.size(); i++) {
            std::string name = std::string("ROI ") + this->dataset->roi[i];
            combo_box->insertSeparator(1000);
            combo_box->addItem(QString::fromStdString(name + " mean"), 100 + i * 4);
            combo_box->addItem(QString::fromStdString(name + " sum"), 100 + i * 4 + 1);
            combo_box->addItem(QString::fromStdString(name + " weighted x"), 100 + i * 4 + 2);
            combo_box->addItem(QString::fromStdString(name + " weighted y"), 100 + i * 4 + 3);
        }
    } else {
        combo_box->clear();
    }
}


void JFJochViewerDatasetInfo::datasetLoaded(std::shared_ptr<const JFJochReaderDataset> dataset) {
    this->dataset = dataset;
    // Per-image mode: the combo holds plot types and the chart follows the displayed image, so a
    // new dataset must not repopulate the combo or run the per-dataset plot.
    if (image_button && image_button->isChecked()) {
        UpdateView();
        return;
    }
    if (dataset) {
        UpdateLabels();
        RestoreSelection();
        UpdatePlot();
    } else {
        chart_view->loadValues({}, 0, false);
        grid_scan_image->clear();
        UpdateLabels();
    }
}

void JFJochViewerDatasetInfo::imageLoaded(std::shared_ptr<const JFJochReaderImage> image) {
    this->image = image;
    if (image) {
        chart_view->setImage(image->ImageData().number);
        grid_scan_image->setImage(image->ImageData().number);
    }
    image_chart->loadImage(image);   // per-image profile follows the displayed image
}

void JFJochViewerDatasetInfo::imageSelectedInChart(int64_t number) {
    emit imageSelected(number, 1);
}

std::vector<float> JFJochViewerDatasetInfo::ExtractMetric(const JFJochReaderDataset &ds, int val,
                                                          bool &one_over_d2) const {
    one_over_d2 = false;
    std::vector<float> data;
    if (val == 0) data = ds.bkg_estimate;
    else if (val == 1) data = ds.spot_count;
    else if (val == 2) data = ds.spot_count_indexed;
    else if (val == 3) data = ds.spot_count_ice_rings;
    else if (val == 4) data = ds.spot_count_low_res;
    else if (val == 5) data = ds.indexing_result;
    else if (val == 6) data = ds.profile_radius;
    else if (val == 7) { data = ds.resolution_estimate; one_over_d2 = true; }
    else if (val == 8) data = ds.b_factor;
    else if (val == 9) data = ds.mosaicity_deg;
    else if (val == 10) data = ds.image_scale_factor;
    else if (val == 11) data = ds.image_scale_cc;
    else if (val == 12) data = ds.integrated_reflections;
    else if (val == 13) data = ds.indexing_lattice_count;
    else if (val == 14) data = ds.ice_ring_ratio;
    else if (val == 15) data = ds.spindle_blind_fraction;
    else if (val == kCompositeMetric) data = ds.protein_score;
    else if (val >= 100) {
        const int roi_index = (val - 100) / 4;
        if (val % 4 == 0) {
            if (roi_index < ds.roi_sum.size() && ds.roi_sum.size() == ds.roi_npixel.size())
                for (int i = 0; i < ds.roi_sum[roi_index].size(); i++)
                    data.push_back(static_cast<float>(ds.roi_sum[roi_index][i])
                                   / static_cast<float>(ds.roi_npixel[roi_index][i]));
        } else if (val % 4 == 1) {
            if (roi_index < ds.roi_sum.size()) {
                data.reserve(ds.roi_sum[roi_index].size());
                for (auto &v: ds.roi_sum[roi_index])
                    data.push_back(v);
            }
        } else if (val % 4 == 2) {
            if (roi_index < ds.roi_x.size()) data = ds.roi_x[roi_index];
        } else if (val % 4 == 3) {
            if (roi_index < ds.roi_y.size()) data = ds.roi_y[roi_index];
        }
    }
    return data;
}

void JFJochViewerDatasetInfo::UpdatePlot() {
    int index = combo_box->currentIndex();
    if (combo_box->count() == 0 || index < 0)
        return;

    int val = combo_box->itemData(index).toInt();

    if (!dataset) {
        grid_button->setEnabled(false);
        return;
    }

    const int64_t image_number = image ? image->ImageData().number : 0;

    bool one_over_d2 = false;
    std::vector<float> data = ExtractMetric(*dataset, val, one_over_d2);

    // One overlay line per other run, plus the in-progress live run. Each run keeps a stable colour
    // tied to its position in the list, so colours don't swap when the active run changes. The
    // primary is matched by dataset pointer (not active_id_) so colour/name stay consistent even
    // while a datasetLoaded/runsChanged pair is still arriving.
    std::vector<JFJochDatasetInfoChartView::NamedSeries> overlays;
    QString primary_name;
    QColor primary_color;
    for (int i = 0; i < runs_.size(); i++) {
        const auto &run = runs_[i];
        if (!run.dataset) continue;
        const QColor color = RunColor(i);
        if (run.dataset == dataset) { primary_name = run.label; primary_color = color; continue; }
        bool ignore = false;
        overlays.push_back({run.label, ExtractMetric(*run.dataset, val, ignore), color, XForRun(*run.dataset)});
    }
    if (live_run_ && live_run_ != dataset) {
        bool ignore = false;
        overlays.push_back({QStringLiteral("Live"), ExtractMetric(*live_run_, val, ignore),
                            RunColor(static_cast<int>(runs_.size())), XForRun(*live_run_)});
    }

    // The whole dataset spans this many images (the largest run, i.e. the original file), so the
    // x-axis is fixed to it and subset runs appear at their real position.
    int64_t full_range = dataset->experiment.GetImageNum();
    for (const auto &run: runs_)
        if (run.dataset)
            full_range = std::max<int64_t>(full_range, run.dataset->experiment.GetImageNum());

    chart_view->loadValues(data, image_number, one_over_d2, dataset.get(), primary_name,
                           std::move(overlays), primary_color, XForRun(*dataset), full_range);

    if (dataset->experiment.GetGridScan()) {
        if (val == kCompositeMetric)
            grid_scan_image->loadComposite(dataset->bkg_estimate, dataset->protein_score,
                                           dataset->ice_score, dataset->grid_crystals,
                                           dataset->experiment.GetGridScan().value());
        else
            grid_scan_image->loadData(data, dataset->experiment.GetGridScan().value(), one_over_d2);
        grid_button->setEnabled(true);
        if (!image_button->isChecked())
            grid_button->setChecked(true);
    } else {
        grid_scan_image->clear();
        grid_button->setEnabled(false);
    }
    UpdateView();
}

void JFJochViewerDatasetInfo::UpdateView() {
    if (image_button->isChecked())
        stack->setCurrentWidget(image_chart);
    else if (grid_button->isChecked() && dataset && dataset->experiment.GetGridScan())
        stack->setCurrentWidget(grid_scan_image);
    else
        stack->setCurrentWidget(chart_view);
}

void JFJochViewerDatasetInfo::runsChanged(QVector<RunData> runs, QString active_id) {
    runs_ = std::move(runs);
    active_id_ = std::move(active_id);
    if (image_button && image_button->isChecked())
        return;   // per-image mode: runs/overlays are a per-dataset concern
    UpdateLabels();
    RestoreSelection();
    UpdatePlot();
}

void JFJochViewerDatasetInfo::liveRunUpdated(std::shared_ptr<const JFJochReaderDataset> in_dataset) {
    live_run_ = std::move(in_dataset);
    if (image_button && image_button->isChecked())
        return;
    UpdatePlot();   // lightweight refresh; no combo rebuild on every live tick
}

void JFJochViewerDatasetInfo::comboBoxSelected(int index) {
    // The one combo means per-dataset metrics in normal mode, per-image plot types in image mode.
    if (image_button && image_button->isChecked()) {
        if (image_chart && index >= 0 && index < combo_box->count())
            image_chart->setPlotType(combo_box->itemData(index).toInt());
    } else {
        UpdatePlot();
    }
}

void JFJochViewerDatasetInfo::setColorMap(int color_map) {
    grid_scan_image->setColorMap(color_map);
}

void JFJochViewerDatasetInfo::resetZoomButtonPressed() {
    if (stack->currentWidget() == grid_scan_image)
        grid_scan_image->fitToView();
    else
        chart_view->resetZoom();
}