Jungfraujoch/viewer/windows/JFJochPixelRefineWindow.cpp

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

#include "JFJochPixelRefineWindow.h"
#include "../image_viewer/JFJochSimpleImage.h"

#include <QWidget>
#include <QHBoxLayout>
#include <QVBoxLayout>
#include <QFormLayout>
#include <QGroupBox>
#include <QFileDialog>
#include <cmath>

JFJochPixelRefineWindow::JFJochPixelRefineWindow(QWidget *parent)
    : JFJochHelperWindow(parent) {
    setWindowTitle("PixelRefine (experimental)");

    auto central = new QWidget(this);
    setCentralWidget(central);
    auto layout = new QHBoxLayout(central);

    // --- predicted image (left, expanding) ---------------------------------
    m_image = new JFJochSimpleImage(this);
    m_image->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Expanding);
    layout->addWidget(m_image, 1);

    // --- control panel (right) ---------------------------------------------
    auto controls = new QWidget(this);
    controls->setMinimumWidth(320);
    auto controlsLayout = new QVBoxLayout(controls);
    layout->addWidget(controls, 0);

    // --- what the left image shows ------------------------------------------
    m_displayMode = new QComboBox(this);
    m_displayMode->addItem(tr("Prediction"));
    m_displayMode->addItem(tr("Squared difference |pred - image|²"));
    m_displayMode->addItem(tr("χ² (weighted residual = LSQ cost)"));
    auto displayForm = new QFormLayout();
    displayForm->addRow(tr("Display:"), m_displayMode);
    controlsLayout->addLayout(displayForm);

    auto paramBox = new QGroupBox(tr("Model parameters"), this);
    auto form = new QFormLayout(paramBox);

    m_R0    = new SliderPlusBox(1e-4, 0.05, 1e-4, 4, this);                       m_R0->setValue(0.005);
    m_R1    = new SliderPlusBox(1e-4, 0.05, 1e-4, 4, this);                       m_R1->setValue(0.005);
    m_bw    = new SliderPlusBox(0.0, 0.05, 1e-4, 4, this);                        m_bw->setValue(0.0);
    m_scale = new SliderPlusBox(1e-3, 1e4, 1e-3, 3, this, SliderPlusBox::Logarithmic); m_scale->setValue(1.0);
    m_B     = new SliderPlusBox(0.0, 200.0, 0.1, 1, this);                        m_B->setValue(0.0);
    m_beamx = new SliderPlusBox(0.0, 5000.0, 0.5, 1, this);                       m_beamx->setValue(0.0);
    m_beamy = new SliderPlusBox(0.0, 5000.0, 0.5, 1, this);                       m_beamy->setValue(0.0);

    form->addRow(tr("R0 radial [Å⁻¹]:"), m_R0);
    form->addRow(tr("R1 tangential [Å⁻¹]:"), m_R1);
    form->addRow(tr("Bandwidth FWHM (Δλ/λ):"), m_bw);
    form->addRow(tr("Scale G:"), m_scale);
    form->addRow(tr("B-factor [Ų]:"), m_B);

    m_overrideBeam = new QCheckBox(tr("Override beam centre"), this);
    form->addRow(QString(), m_overrideBeam);
    form->addRow(tr("Beam X [px]:"), m_beamx);
    form->addRow(tr("Beam Y [px]:"), m_beamy);
    m_beamx->setEnabled(false);
    m_beamy->setEnabled(false);

    controlsLayout->addWidget(paramBox);

    // --- what "Refine" is allowed to move ----------------------------------
    auto refBox = new QGroupBox(tr("Refine (Ceres)"), this);
    auto refLayout = new QVBoxLayout(refBox);
    m_refOrientation = new QCheckBox(tr("Orientation"), this);   m_refOrientation->setChecked(true);
    m_refCell        = new QCheckBox(tr("Unit cell"), this);
    m_refBeam        = new QCheckBox(tr("Beam centre"), this);
    m_refScale       = new QCheckBox(tr("Scale G"), this);       m_refScale->setChecked(true);
    m_refB           = new QCheckBox(tr("B-factor"), this);
    m_refR           = new QCheckBox(tr("Widths R0/R1"), this);  m_refR->setChecked(true);
    for (auto *cb : {m_refOrientation, m_refCell, m_refBeam, m_refScale, m_refB, m_refR})
        refLayout->addWidget(cb);
    controlsLayout->addWidget(refBox);

    // --- buttons + readouts -------------------------------------------------
    m_loadRef = new QPushButton(tr("Load reference MTZ…"), this);
    m_refine  = new QPushButton(tr("Refine"), this);
    controlsLayout->addWidget(m_loadRef);
    controlsLayout->addWidget(m_refine);

    m_residual = new QLabel(tr("Residual: —"), this);
    m_pipelineCC = new QLabel(tr("Pipeline CC (ref): —"), this);
    m_status   = new QLabel(QString(), this);
    m_status->setWordWrap(true);
    m_status->setStyleSheet("color: rgb(80, 80, 80);");
    controlsLayout->addWidget(m_residual);
    controlsLayout->addWidget(m_pipelineCC);
    controlsLayout->addWidget(m_status);
    controlsLayout->addStretch(1);

    // --- debounce timer for live preview -----------------------------------
    m_debounce = new QTimer(this);
    m_debounce->setSingleShot(true);
    m_debounce->setInterval(150);
    connect(m_debounce, &QTimer::timeout, this, [this] {
        emit paramsChanged(currentParams());
    });

    for (auto *s : {m_R0, m_R1, m_bw, m_scale, m_B, m_beamx, m_beamy})
        connect(s, &SliderPlusBox::valueChanged, this, [this](double) { onControlChanged(); });

    connect(m_displayMode, &QComboBox::currentIndexChanged, this, [this](int) { onControlChanged(); });

    connect(m_overrideBeam, &QCheckBox::toggled, this, [this](bool on) {
        m_beamx->setEnabled(on);
        m_beamy->setEnabled(on);
        onControlChanged();
    });

    connect(m_loadRef, &QPushButton::clicked, this, [this] {
        const QString path = QFileDialog::getOpenFileName(
            this, tr("Load reference MTZ"), QString(), tr("MTZ files (*.mtz);;All files (*)"));
        if (!path.isEmpty())
            emit loadReferenceRequested(path);
    });

    connect(m_refine, &QPushButton::clicked, this, [this] {
        // Cancel any pending live-preview: otherwise a debounce armed by a slider
        // move just before this click fires after the refine and overwrites the
        // refined residual/preview with the stale pre-refine slider values.
        m_debounce->stop();
        PixelRefineParams p = currentParams();
        p.max_iterations = 5;
        emit refineRequested(p);
    });
}

PixelRefineParams JFJochPixelRefineWindow::currentParams() const {
    PixelRefineParams p;
    p.R0 = m_R0->value();
    p.R1 = m_R1->value();
    p.bandwidth_fwhm = m_bw->value();
    p.scale_factor = m_scale->value();
    p.B_factor = m_B->value();
    if (m_overrideBeam->isChecked()) {
        p.beam_x = m_beamx->value();
        p.beam_y = m_beamy->value();
    } else {
        p.beam_x = NAN;
        p.beam_y = NAN;
    }
    p.refine_orientation = m_refOrientation->isChecked();
    p.refine_unit_cell   = m_refCell->isChecked();
    p.refine_beam_center = m_refBeam->isChecked();
    p.refine_scale       = m_refScale->isChecked();
    p.refine_B           = m_refB->isChecked();
    p.refine_R           = m_refR->isChecked();
    p.display_mode       = m_displayMode->currentIndex();
    return p;
}

void JFJochPixelRefineWindow::onControlChanged() {
    if (m_suppress)
        return;
    m_debounce->start();
}

void JFJochPixelRefineWindow::imageLoaded(std::shared_ptr<const JFJochReaderImage> image) {
    if (!image)
        return;

    // Initialise the beam-centre sliders from the geometry once.
    if (!m_beamInit) {
        const auto geom = image->Dataset().experiment.GetDiffractionGeometry();
        m_beamx->setMax(static_cast<double>(image->Dataset().experiment.GetXPixelsNum()));
        m_beamy->setMax(static_cast<double>(image->Dataset().experiment.GetYPixelsNum()));
        m_suppress = true;
        m_beamx->setValue(geom.GetBeamX_pxl());
        m_beamy->setValue(geom.GetBeamY_pxl());
        m_suppress = false;
        m_beamInit = true;
    }

    // Show the standard ScaleOnTheFly pipeline's per-image CC vs the reference (set
    // on the message during analysis when a reference is loaded), as a baseline to
    // compare PixelRefine against.
    const auto pipeline_cc = image->ImageData().image_scale_cc;
    if (pipeline_cc.has_value() && std::isfinite(pipeline_cc.value()))
        m_pipelineCC->setText(tr("Pipeline CC (ref): %1%")
                                  .arg(pipeline_cc.value() * 100.0, 0, 'f', 1));
    else
        m_pipelineCC->setText(tr("Pipeline CC (ref): —"));

    // Request a predicted-image preview for the (re)loaded image. Without this the
    // preview only refreshed on a slider change, so opening/reanalyzing an image
    // left the predicted view empty. The worker no-ops if there is no reference or
    // the image is not integrated yet.
    onControlChanged();
}

void JFJochPixelRefineWindow::setPredictedImage(std::shared_ptr<const SimpleImage> image) {
    m_image->setImage(std::move(image));
}

void JFJochPixelRefineWindow::setResidual(double cost, double cc, int64_t n_reflections) {
    const QString cc_str = std::isfinite(cc) ? QString::number(cc * 100.0, 'f', 1) + "%"
                                             : QStringLiteral("—");
    m_residual->setText(tr("Residual: %1   CC: %2   (%3 reflections)")
                            .arg(cost, 0, 'g', 6)
                            .arg(cc_str)
                            .arg(n_reflections));
}

void JFJochPixelRefineWindow::setRefinedParams(PixelRefineParams params) {
    m_suppress = true;
    m_R0->setValue(params.R0);
    m_R1->setValue(params.R1);
    m_bw->setValue(params.bandwidth_fwhm);
    m_scale->setValue(params.scale_factor);
    m_B->setValue(params.B_factor);
    if (std::isfinite(params.beam_x) && std::isfinite(params.beam_y)) {
        m_beamx->setValue(params.beam_x);
        m_beamy->setValue(params.beam_y);
    }
    m_suppress = false;
}

void JFJochPixelRefineWindow::setStatus(QString message) {
    m_status->setText(message);
}