// 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;
    R = static_cast<int>(std::ceil(r2));
    G = 2 * R + 1;
    GG = G * G;

    // The X-ray bandwidth enters ONE place: it smears a reflection radially by bw_sigma * Rpx, which
    // the per-reflection Gaussian carries as part of its radial variance. It is not a mode switch -
    // the background estimator and the parallax/capture term below are the same whatever the beam is.
    bw_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f;

    const double c_par = parallax_var_px2(det.GetSensorMaterial(), det.GetSensorThickness_um(),
                                          geom.GetWavelength_A(), geom.GetPixelSize_mm() * 1000.0);
    c_radial = c_par + 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;

    // Per-reflection signal/background geometry: the ring elongated radially by k_sigma times the
    // beam's own radial streak, capped. k_sigma = 0 is the fixed circular stencil, bit for bit, and
    // so is any monochromatic beam, where the streak is zero.
    stencil.beam_x = beam_x;
    stencil.beam_y = beam_y;
    stencil.r2 = r2;
    stencil.r3 = r3;
    stencil.bw_sigma = static_cast<float>(bw_sigma);
    stencil.k_sigma = settings.GetStencilKSigma();
    stencil.max_grow = bragg_engine::MAX_STENCIL_GROW_OVER_R3 * r3;

    // Robust background ring, one estimator or the other (see BraggIntegrationSettings): a high-side
    // sigma-clip (rugnux --background-clip, the default) or, when the clip is switched off, a
    // symmetric trimmed mean (rugnux --background-trim). The caller owns the choice - the engine no
    // longer overrides it for broadband data.
    bkg_clip_nsigma = settings.GetBackgroundClipNSigma();
    bkg_trim = bkg_clip_nsigma > 0.0f ? 0.0f : settings.GetBackgroundTrimFraction();

    // Overlap treatment. Ownership is decided out to the fit grid's half size, which is where the
    // profile fit reads pixels; beyond it a pixel that nobody claims is this reflection's own.
    // Excluding the shared pixels needs a profile to renormalise, so it cannot act on a box sum -
    // drop it to Off there rather than build an owner map nothing will read.
    overlap = settings.GetOverlap();
    if (overlap == OverlapMode::Exclude && mode == IntegratorMode::BoxSum)
        overlap = OverlapMode::Off;
    overlap_min_peak = settings.GetOverlapMinPeak();
    claim = static_cast<float>(R);
    inv_claim = 1.0f / claim;

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

    // The table spans zero growth up to whatever the widest reflection on this detector reaches, one
    // kernel per pixel of growth; with nothing elongated a single kernel is all there is, which is
    // the layout and the values of every build before the stencil existed. It is built only when the
    // correction can ever run - the rows are not cheap, and nothing may read them otherwise:
    // bkg_radial is raised after construction only by the auto mode (MXAnalysisWithoutFPGA), which
    // requires bkg_radial_auto, and the GPU allocates its curve buffers under the same condition.
    // n_kern is the largest row BraggStencilKernelIndex can select, plus one.
    r_max = std::hypot(std::max<double>(beam_x, static_cast<double>(xpixel) - beam_x),
                       std::max<double>(beam_y, static_cast<double>(ypixel) - beam_y));
    bkg_radial_built = bkg_radial || bkg_radial_auto;
    const float grow_max = bkg_radial_built ? BraggStencilGrow_px(static_cast<float>(r_max), stencil)
                                            : 0.0f;
    n_kern = static_cast<int>(std::lround(grow_max)) + 1;
    // Every row must fit: the last one is built at grow = n_kern - 1, which rounding can put just
    // above grow_max.
    k_off = static_cast<int>(std::ceil(r3 + std::max<double>(grow_max, n_kern - 1))) + 1;
    k_len = 2 * k_off + 1;
    k_diff.clear();
    k_diff.reserve(static_cast<size_t>(n_kern) * k_len);
    for (int j = 0; j < n_kern; ++j)
        BuildRadialKernel(static_cast<float>(j));

    polarization = experiment.GetPolarizationFactor();
}

void BraggIntegrationEngine::BuildRadialKernel(float grow) {
    // Histogram the stencil over radial offset, averaged over azimuth so the kernel does not depend
    // on where the reflection sits. The average is over the SUB-PIXEL PHASE of the detector grid
    // against the radial direction, not over the stencil's own orientation: the stencil is built in
    // the reflection's frame at each azimuth, so an elongated one stays aligned with the radius, as
    // it is on the detector. 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).
    // The signal disk is a circle whatever the ring does, so its histogram is the same for every
    // kernel in the table - build it once.
    const bool first = hist_disk.empty();
    if (first)
        hist_disk.assign(k_len, 0.0);
    std::vector<double> hist_ann(k_len, 0.0);
    constexpr int n_phi = 512;
    const int span = static_cast<int>(std::ceil(r3 + grow)) + 1;
    const float si = r2 / (r2 + grow), so = r3 / (r3 + grow);
    const double q_in = 1.0 - static_cast<double>(si) * si;
    const double q_out = 1.0 - static_cast<double>(so) * so;
    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 double rad = dx * cp + dy * sp;
                const int k = k_off + static_cast<int>(std::lround(rad));
                if (k < 0 || k >= k_len)
                    continue;
                const double rad2 = rad * rad;
                if (d2 < r1_sq) {
                    if (first) hist_disk[k] += 1.0;
                } else if (d2 - q_in * rad2 >= r2_sq && d2 - q_out * rad2 < r3_sq) {
                    hist_ann[k] += 1.0;
                }
            }
    }
    if (first) sum_disk = std::accumulate(hist_disk.begin(), hist_disk.end(), 0.0);
    const double sd = sum_disk;
    const double sa = std::accumulate(hist_ann.begin(), hist_ann.end(), 0.0);
    for (int k = 0; k < k_len; ++k)
        k_diff.push_back(static_cast<float>(hist_ann[k] / sa - hist_disk[k] / sd));
}

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;
        refl.var_bkg = fr.var_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;
}