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

#include "BraggIntegrationEngine.h"

#include <algorithm>
#include <string>

namespace {

// Radial parallax broadening as the coefficient of tan^2(2theta), i.e. Var(z)/pixel^2 [px^2].
// Copied verbatim from ProfileIntegrate2D: a photon converts at a random depth z (exponential,
// attenuation length L, truncated at the sensor thickness), shifting the recorded spot radially by
// z*tan(2theta). L is photoelectric-dominated (~lambda^3), so a per-material reference (13 keV) is
// scaled by lambda^3; Si and CdTe are the sensors in use.
double parallax_var_px2(const std::string &material, double thickness_um, double lambda_A, double pixel_um) {
    if (!(thickness_um > 0.0) || !(pixel_um > 0.0) || !(lambda_A > 0.0))
        return 0.0;
    const double L_ref = material == "CdTe" ? 42.6 : 273.0;   // attenuation length [um] at 0.953 A
    const double s = lambda_A / 0.953;
    const double L = L_ref / (s * s * s);
    const double a = thickness_um / L, e = std::exp(-a);
    if (1.0 - e <= 0.0)
        return 0.0;
    const double mean = L * (1.0 - (1.0 + a) * e) / (1.0 - e);
    const double ez2 = L * L * (2.0 - (a * a + 2.0 * a + 2.0) * e) / (1.0 - e);
    const double var = std::max(0.0, ez2 - mean * mean);     // um^2
    return var / (pixel_um * pixel_um);
}

}   // namespace

BraggIntegrationEngine::BraggIntegrationEngine(const DiffractionExperiment &experiment)
    : geom(experiment.GetDiffractionGeometry()) {
    const auto settings = experiment.GetBraggIntegrationSettings();
    const auto &det = experiment.GetDetectorSetup();

    mode = settings.GetIntegrator();
    empirical = mode == IntegratorMode::ProfileEmpirical;

    // Same frame as the reflections' predicted_x/predicted_y and the ImagePreprocessorBuffer that
    // feeds this engine (MXAnalysisWithoutFPGA sizes that buffer to GetPixelsNum()).
    xpixel = experiment.GetXPixelsNum();
    ypixel = experiment.GetYPixelsNum();
    npixel = experiment.GetPixelsNum();

    r1_sq = settings.GetR1() * settings.GetR1();
    r2 = settings.GetR2();
    r2_sq = r2 * r2;
    r3 = settings.GetR3();
    r3_sq = r3 * r3;
    min_sigma_ratio = settings.GetMinimumSigmaInRegardsToI();
    R = static_cast<int>(std::ceil(r2));
    G = 2 * R + 1;
    GG = G * G;

    // A set bandwidth (broadband / stills) vs monochromatic (rotation) splits the treatment: the
    // background sigma-clip and radial-elongation terms are path-dependent (see ProfileIntegrate2D).
    bw_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f;
    broadband = bw_sigma > 0.0;
    apply_bkg_clip = broadband;

    const double c_par = parallax_var_px2(det.GetSensorMaterial(), det.GetSensorThickness_um(),
                                          geom.GetWavelength_A(), geom.GetPixelSize_mm() * 1000.0);
    c_radial = c_par + (broadband ? 0.0 : bragg_engine::C_CAPTURE);
    F_px = geom.GetDetectorDistance_mm() / std::max(1e-6f, geom.GetPixelSize_mm());
    beam_x = geom.GetBeamX_pxl();
    beam_y = geom.GetBeamY_pxl();
    use_ellipse = !empirical && (bw_sigma > 0.0 || c_radial > 0.0);

    polarization = experiment.GetPolarizationFactor();
}

std::vector<Reflection> BraggIntegrationEngine::Finalize(const std::vector<Reflection> &predicted,
                                                         size_t npredicted,
                                                         const std::vector<BraggFitResult> &results,
                                                         int64_t image_number) const {
    std::vector<Reflection> out;
    out.reserve(npredicted);
    for (size_t i = 0; i < npredicted; ++i) {
        const auto &fr = results[i];
        if (!fr.ok)
            continue;
        Reflection refl = predicted[i];
        refl.I = fr.I;
        refl.sigma = fr.sigma;
        refl.bkg = fr.bkg;
        if (fr.has_observed) {
            refl.observed_x = fr.observed_x;
            refl.observed_y = fr.observed_y;
        }
        refl.observed = true;
        if (polarization)
            refl.rlp /= geom.CalcAzIntPolarizationCorr(refl.predicted_x, refl.predicted_y, polarization.value());
        refl.image_scale_corr = refl.rlp / refl.partiality;
        refl.image_number = static_cast<float>(image_number);
        out.push_back(refl);
    }
    return out;
}