// 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 <QSettings>

#include "JFJochViewerDatasetInfo.h"

namespace {
    // Combo data value of the combined spots + background plot.
    constexpr int SPOTS_PLUS_BACKGROUND = 16;

    // Grid scans and standard scans keep separate saved favourites.
    QString FavouriteGroup(bool grid_scan) {
        return grid_scan ? QStringLiteral("datasetInfoPlot/gridScan")
                         : QStringLiteral("datasetInfoPlot/standard");
    }

    // 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];
    }

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

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

    reset_button = new QPushButton("Reset zoom", this);

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

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

    add_plot_button = new QPushButton("+ Plot", this);
    UpdateButtonWidths();
    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_view2 = new JFJochDatasetInfoChartView(this);
    chart_view2->hide();                 // only "Spots + background" shows the second plot
    chart_view2->SetMinYZero(false);     // the background plot is always about the shape, not the level
    chart_view->SetCompactPeer(chart_view2);
    chart_view2->SetCompactPeer(chart_view);
    grid_scan_image = new JFJochGridScanImage(this);
    image_chart = new JFJochViewerSidePanelChart(this);   // per-image profiles, folded in from the side panel

    charts_holder = new QWidget(this);
    auto *charts_layout = new QVBoxLayout(charts_holder);
    charts_layout->setContentsMargins(0, 0, 0, 0);
    charts_layout->setSpacing(0);
    charts_layout->addWidget(chart_view, 1);
    charts_layout->addWidget(chart_view2, 1);   // equal stretch: the two stacked plots share the height

    stack->addWidget(charts_holder);
    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(chart_view2, &JFJochDatasetInfoChartView::imageSelected,
            this, &JFJochViewerDatasetInfo::imageSelectedInChart);

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

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

    connect(combo_box, &QComboBox::currentIndexChanged,
            this, &JFJochViewerDatasetInfo::comboBoxSelected);
    // activated (not currentIndexChanged) fires only on a user pick, which is exactly when the
    // live default must stop overriding the selection - and when the saved favourite follows it.
    connect(combo_box, &QComboBox::activated, this, [this](int) {
        user_picked_ = true;
        (combo_is_image_mode_ ? last_image_selection : last_selection) = combo_box->currentData().toInt();
        SaveFavourite();
    });

    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);   // leaves per-image mode: the combo must follow
                ModeChanged();
                SaveFavourite();
            });
    connect(image_button, &QPushButton::clicked,
            [this]() {
                if (image_button->isChecked())
                    grid_button->setChecked(false);
                ModeChanged();
                SaveFavourite();
            });
    connect(image_chart, &JFJochViewerSidePanelChart::writeStatusBar,
            [this](QString s, int t) { emit writeStatusBar(s, t); });
    setLayout(layout);
}

void JFJochViewerDatasetInfo::RestoreSelection() {
    QSignalBlocker bl(combo_box);
    // On a live source, spots against background is what says whether beam and crystal are still
    // delivering, so that is where a live session starts - until the user picks something else.
    const int live_default = combo_box->findData(SPOTS_PLUS_BACKGROUND);
    if (!combo_is_image_mode_ && http_connected_ && !user_picked_ && live_default >= 0) {
        combo_box->setCurrentIndex(live_default);
        return;
    }
    // Each mode remembers its own selection by the item's data code: the two lists are unrelated,
    // and a code survives the combo being rebuilt with a different item set - one that is no
    // longer offered simply isn't found and the first item applies.
    const int index = combo_box->findData(combo_is_image_mode_ ? last_image_selection : last_selection);
    combo_box->setCurrentIndex(index >= 0 ? index : 0);
}

void JFJochViewerDatasetInfo::SaveFavourite() {
    if (!dataset)
        return;
    QSettings settings("PSI", "jfjoch_viewer");
    settings.beginGroup(FavouriteGroup(dataset->experiment.GetGridScan().has_value()));
    settings.setValue("metric", last_selection);
    settings.setValue("imagePlot", last_image_selection);
    settings.setValue("perImage", image_button->isChecked());
    settings.setValue("grid", grid_button->isChecked());
}

void JFJochViewerDatasetInfo::resetPlotDefaults() {
    user_picked_ = false;
    image_button->setChecked(false);
    // A grid scan starts on its 2D map, like a first load with nothing saved.
    grid_button->setChecked(dataset && dataset->experiment.GetGridScan().has_value());
    UpdateLabels();       // rebuild the combo for per-dataset mode
    last_selection = 0;   // default selections, after UpdateLabels captured the old ones
    last_image_selection = 0;
    RestoreSelection();
    UpdatePlot();
    UpdateView();
}

void JFJochViewerDatasetInfo::ApplyFavourite(bool grid_scan) {
    QSettings settings("PSI", "jfjoch_viewer");
    settings.beginGroup(FavouriteGroup(grid_scan));
    image_button->setChecked(settings.value("perImage", false).toBool());
    // With nothing saved, a grid scan starts on its 2D map.
    grid_button->setChecked(settings.value("grid", grid_scan).toBool());
    // ModeChanged's sequence, with the saved codes injected after UpdateLabels has captured the
    // outgoing selection - the other way round the capture would overwrite them.
    UpdateLabels();
    last_selection = settings.value("metric", last_selection).toInt();
    last_image_selection = settings.value("imagePlot", last_image_selection).toInt();
    RestoreSelection();
    if (combo_is_image_mode_)
        comboBoxSelected(combo_box->currentIndex());
    else
        UpdatePlot();
    UpdateView();
}

void JFJochViewerDatasetInfo::ModeChanged() {
    UpdateLabels();          // swap the one combo's contents for the mode
    RestoreSelection();
    if (combo_is_image_mode_)
        comboBoxSelected(combo_box->currentIndex());   // apply the per-image plot type
    else
        UpdatePlot();
    UpdateView();
}

void JFJochViewerDatasetInfo::setHttpConnection(bool connected, QString) {
    http_connected_ = connected;
    if (!connected || user_picked_ || (image_button && image_button->isChecked()))
        return;
    RestoreSelection();   // the combo may still be empty; datasetLoaded() re-runs this either way
    UpdatePlot();
}

// The buttons share one width so the row stays even. It is a character count rather than a pixel
// count, so the labels still fit when the font grows (102 px at the default font, as good as the
// 100 px it replaces).
void JFJochViewerDatasetInfo::UpdateButtonWidths() {
    const int char_width = fontMetrics().averageCharWidth();
    reset_button->setFixedWidth(17 * char_width);
    grid_button->setFixedWidth(17 * char_width);
    image_button->setFixedWidth(17 * char_width);
    add_plot_button->setFixedWidth(13 * char_width);
}

void JFJochViewerDatasetInfo::changeEvent(QEvent *event) {
    QWidget::changeEvent(event);
    if (event->type() == QEvent::FontChange)
        UpdateButtonWidths();
}

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

    if (combo_box->count() > 0)
        (combo_is_image_mode_ ? last_image_selection : last_selection) = combo_box->currentData().toInt();

    combo_box->clear();
    combo_is_image_mode_ = image_button && image_button->isChecked();

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

    if (this->dataset) {
        combo_box->addItem("Background estimate", 0);
        combo_box->addItem("Spots + background", SPOTS_PLUS_BACKGROUND);
        combo_box->addItem("Ice ring score", 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;
    // The first dataset of a kind (grid scan / standard) starts from the saved favourite for that
    // kind - the user's last explicit plot choice and toggle states on such a dataset.
    if (dataset) {
        const bool grid_scan = dataset->experiment.GetGridScan().has_value();
        if (favourite_kind_ != grid_scan) {
            favourite_kind_ = grid_scan;
            ApplyFavourite(grid_scan);   // rebuilds the combo for the favourite's mode and re-plots
            return;
        }
    }
    // 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);
        chart_view2->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_score;
    else if (val == 15) data = ds.spindle_blind_fraction;
    else if (val == SPOTS_PLUS_BACKGROUND) data = ds.spot_count;   // background is the secondary line, see UpdatePlot
    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_)});
    }

    // "Spots + background": two half-height plots stacked in the same slot - the spot count on
    // top (its x labels hidden), the background below, carrying the x axis for both.
    const bool split = (val == SPOTS_PLUS_BACKGROUND);
    if (split && !primary_color.isValid())
        primary_color = RunColor(0);
    chart_view->SetCompact(split ? QStringLiteral("Spots") : QString(), split);
    chart_view2->SetCompact(split ? QStringLiteral("Bkg") : QString(), false);
    chart_view2->setVisible(split);

    // 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 (split) {
        chart_view2->loadValues(dataset->bkg_estimate, image_number, false, dataset.get(),
                                QStringLiteral("Background"), {}, RunColor(1), XForRun(*dataset),
                                full_range);
    }

    if (dataset->experiment.GetGridScan()) {
        grid_scan_image->loadData(data, dataset->experiment.GetGridScan().value(), one_over_d2);
        grid_button->setEnabled(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(charts_holder);
}

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