// 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 <cmath>
#include <numeric>
#include <string>

#include "../../common/JFJochMath.h"   // PI (M_PI is not standard, and MSVC does not define it)

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;

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

    // Robust background ring, one estimator or the other (see BraggIntegrationSettings). Broadband
    // (non-zero bandwidth: pink-beam / DMM) data keep their tuned 3 sigma high-side clip whatever the
    // settings say; monochromatic data - rotation AND stills, the discriminator is the beam, not the
    // acquisition mode - take the clip multiplier from settings, and fall back to the symmetric trim
    // only when the clip is switched off (rugnux --background-trim).
    bkg_clip_nsigma = broadband ? 3.0f : settings.GetBackgroundClipNSigma();
    bkg_trim = (broadband || bkg_clip_nsigma > 0.0f) ? 0.0f : settings.GetBackgroundTrimFraction();

    // Radial-offset kernels for the background curvature correction. A stencil pixel at (dx, dy)
    // sits at radial offset dx*cos(phi) + dy*sin(phi) from the reflection, where phi is the
    // reflection's azimuth; averaging over phi makes the kernels position-independent, which is
    // exact to the extent the stencil is small against the reflection's radius (r3 = 10 px vs
    // hundreds). k_diff is the annulus histogram minus the disk histogram, each normalised, so
    // dot(k_diff, B) is directly mean_annulus(B) - mean_disk(B).
    // Unset = auto: start off, and let the analysis raise it per image where the ice score says the
    // background really is radial. An engine nobody drives therefore never applies the correction.
    const auto radial = settings.GetBackgroundRadialCorrection();
    bkg_radial_auto = !radial.has_value();
    bkg_radial = radial.value_or(false);
    k_off = static_cast<int>(std::ceil(r3)) + 1;
    k_diff.assign(2 * k_off + 1, 0.0f);
    {
        std::vector<double> hist_disk(k_diff.size(), 0.0), hist_ann(k_diff.size(), 0.0);
        constexpr int n_phi = 512;
        const int span = static_cast<int>(std::ceil(r3)) + 1;
        for (int p = 0; p < n_phi; ++p) {
            const double phi = 2.0 * PI * p / n_phi, cp = std::cos(phi), sp = std::sin(phi);
            for (int dy = -span; dy <= span; ++dy)
                for (int dx = -span; dx <= span; ++dx) {
                    const double d2 = static_cast<double>(dx) * dx + static_cast<double>(dy) * dy;
                    const int k = k_off + static_cast<int>(std::lround(dx * cp + dy * sp));
                    if (k < 0 || k >= static_cast<int>(k_diff.size()))
                        continue;
                    if (d2 < r1_sq) hist_disk[k] += 1.0;
                    else if (d2 >= r2_sq && d2 < r3_sq) hist_ann[k] += 1.0;
                }
        }
        const double sd = std::accumulate(hist_disk.begin(), hist_disk.end(), 0.0);
        const double sa = std::accumulate(hist_ann.begin(), hist_ann.end(), 0.0);
        for (size_t k = 0; k < k_diff.size(); ++k)
            k_diff[k] = static_cast<float>(hist_ann[k] / sa - hist_disk[k] / sd);
    }

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