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

#include <QGridLayout>
#include <QPushButton>

#include "JFJochViewerDatasetInfo.h"

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

    layout->addWidget(reset_button, 0, 2);
    layout->addWidget(grid_button, 0, 3);

    stack = new QStackedWidget(this);
    chart_view = new JFJochDatasetInfoChartView(this);
    grid_scan_image = new JFJochGridScanImage(this);

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

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

    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() && dataset && dataset->experiment.GetGridScan())
                    stack->setCurrentWidget(grid_scan_image);
                else
                    stack->setCurrentWidget(chart_view);
            });
    setLayout(layout);
}

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

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

    combo_box->clear();

    if (this->dataset) {
        combo_box->addItem("Background estimate", 0);
        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);

        if (!dataset->indexing_result.empty()) {
            combo_box->insertSeparator(1000);
            combo_box->addItem("Indexing result", 5);
            combo_box->addItem("Profile radius", 6);
            combo_box->addItem("B-factor", 8);
        }

        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;
    if (dataset) {
        UpdateLabels();
        if (last_selection < combo_box->count())
            combo_box->setCurrentIndex(last_selection);
        else
            combo_box->setCurrentIndex(0);
        UpdatePlot();
    } else {
        chart_view->loadValues<float>({}, 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);
    }
}

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

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) {
        int64_t image_number = 0;
        if (image)
            image_number = image->ImageData().number;

        bool one_over_d2 = false;
        std::vector<float> data;
        if (val == 0)
            data = dataset->bkg_estimate;
        else if (val == 1)
            data = dataset->spot_count;
        else if (val == 2)
            data = dataset->spot_count_indexed;
        else if (val == 3)
            data = dataset->spot_count_ice_rings;
        else if (val == 4)
            data = dataset->spot_count_low_res;
        else if (val == 5)
            data = dataset->indexing_result;
        else if (val == 6)
            data = dataset->profile_radius;
        else if (val == 7) {
            data = dataset->resolution_estimate;
            one_over_d2 = true;
        } else if (val == 8) {
            data = dataset->b_factor;
        } else if (val >= 100) {
            int roi_index = (val - 100) / 4;

            if (val % 4 == 0) {
                if (roi_index < dataset->roi_sum.size() && dataset->roi_sum.size() == dataset->roi_npixel.size()) {
                    for(int i = 0; i < dataset->roi_sum[roi_index].size(); i++) {
                        data.push_back(static_cast<float>(dataset->roi_sum[roi_index][i])
                            / static_cast<float>(dataset->roi_npixel[roi_index][i]));
                    }
                }
            } else if (val % 4 == 1) {
                if (roi_index < dataset->roi_sum.size()) {
                    data.reserve(dataset->roi_sum[roi_index].size());
                    for (auto &v: dataset->roi_sum[roi_index])
                        data.push_back(v);
                }
            } else if (val % 4 == 2) {
                if (roi_index < dataset->roi_x.size())
                    data = dataset->roi_x[roi_index];
            } else if (val % 4 == 3) {
                if (roi_index < dataset->roi_y.size())
                    data = dataset->roi_y[roi_index];
            }
        }

        chart_view->loadValues(data, image_number, one_over_d2, dataset.get());
        if (dataset->experiment.GetGridScan()) {
            stack->setCurrentWidget(grid_scan_image);
            grid_scan_image->loadData(data, dataset->experiment.GetGridScan().value(), one_over_d2);
            grid_button->setEnabled(true);
            grid_button->setChecked(true);
        } else {
            grid_scan_image->clear();
            stack->setCurrentWidget(chart_view);
            grid_button->setEnabled(false);
        }
    } else {
        grid_button->setEnabled(false);
    }
}

void JFJochViewerDatasetInfo::comboBoxSelected(int index) {
    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();
}