Jungfraujoch/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp

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

#include <algorithm>
#include <bitset>
#include <cmath>

#include "AdaptiveSpotFinderCPU.h"
#include "AdaptiveThreshold.h"

AdaptiveSpotFinderCPU::AdaptiveSpotFinderCPU(const AzimuthalIntegrationMapping &in_mapping)
    : ImageSpotFinderCPU(static_cast<int32_t>(in_mapping.GetWidth()),
                         static_cast<int32_t>(in_mapping.GetHeight())),
      mapping(in_mapping) {
    const size_t nbins = mapping.GetBinNumber();
    ring_sum.assign(nbins, 0);
    ring_sum2.assign(nbins, 0);
    ring_cnt.assign(nbins, 0);
    ring_mean.assign(nbins, 0.0f);
    ring_sigma.assign(nbins, 0.0f);
    ring_thr.assign(nbins, 0.0f);
    ring_bkg.assign(nbins, NAN);
    ring_bits.assign(OutputSize(), 0);
}

// Accumulate per-ring mean/variance from the raw (photon) image. clip_k <= 0 -> use every valid
// pixel (first pass); clip_k > 0 -> keep only pixels within clip_k sigma of the current ring mean,
// which removes the Bragg peaks from the background estimate.
void AdaptiveSpotFinderCPU::AccumulateRings(const ImagePreprocessorBuffer &image, float clip_k) {
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const size_t nbins = ring_sum.size();
    const size_t npix = static_cast<size_t>(width) * height;

    std::fill(ring_sum.begin(), ring_sum.end(), 0);
    std::fill(ring_sum2.begin(), ring_sum2.end(), 0);
    std::fill(ring_cnt.begin(), ring_cnt.end(), 0);

    for (size_t pxl = 0; pxl < npix; ++pxl) {
        const int32_t v = image[pxl];
        if (v == INT32_MIN || v == INT32_MAX) continue;   // bad / saturated
        const uint16_t b = pixel_to_bin[pxl];
        if (b >= nbins) continue;                          // masked / out of range (UINT16_MAX)
        if (clip_k > 0.0f) {
            const float lo = ring_mean[b] - clip_k * ring_sigma[b];
            const float hi = ring_mean[b] + clip_k * ring_sigma[b];
            if (v < lo || v > hi) continue;                // exclude peaks / outliers
        }
        ring_sum[b] += v;
        ring_sum2[b] += static_cast<uint64_t>(static_cast<int64_t>(v) * v);
        ring_cnt[b] += 1;
    }

    for (size_t b = 0; b < nbins; ++b) {
        if (ring_cnt[b] > 0) {
            const double m = static_cast<double>(ring_sum[b]) / ring_cnt[b];
            const double var = std::max(0.0, static_cast<double>(ring_sum2[b]) / ring_cnt[b] - m * m);
            ring_mean[b] = static_cast<float>(m);
            ring_sigma[b] = static_cast<float>(std::sqrt(var));
        }
    }
}

void AdaptiveSpotFinderCPU::Detect(const ImagePreprocessorBuffer &image,
                                   const SpotFindingSettings &settings) {
    const size_t nbins = ring_sum.size();

    // --- Stage A: robust per-ring background (one plain pass + two sigma-clip passes) ---
    AccumulateRings(image, 0.0f);
    AccumulateRings(image, 3.0f);
    AccumulateRings(image, 3.0f);

    // --- Stage B: per-ring threshold from the single portable knob E (false pixels / frame) ---
    int64_t n_total = 0;
    double g_sum = 0.0, g_sum2 = 0.0;
    for (size_t b = 0; b < nbins; ++b) {
        n_total += ring_cnt[b];
        g_sum += static_cast<double>(ring_sum[b]);
        g_sum2 += static_cast<double>(ring_sum2[b]);
    }
    if (n_total == 0) {
        // Nothing valid to threshold against: leave no strong pixels for ExtractSpots to build on.
        std::fill(output_buffer.begin(), output_buffer.end(), 0);
        std::fill(ring_bkg.begin(), ring_bkg.end(), NAN);
        return;
    }

    for (size_t b = 0; b < nbins; ++b)
        ring_bkg[b] = (ring_cnt[b] < adaptive_threshold::MIN_RING_PIXELS) ? NAN : ring_mean[b];

    const double E = std::max(1.0f, settings.false_pixels_per_frame);
    double p = E / static_cast<double>(n_total);
    p = std::min(std::max(p, 1e-9), 0.1);
    const float z = static_cast<float>(adaptive_threshold::NormalQuantile(1.0 - p));

    // whole-frame fallback background for rings too sparse to trust on their own
    const double g_mean = g_sum / n_total;
    const double g_sigma = std::sqrt(std::max(0.0, g_sum2 / n_total - g_mean * g_mean));
    const float g_thr = adaptive_threshold::RingThreshold(static_cast<float>(g_mean),
                                                          static_cast<float>(g_sigma), p, z);

    for (size_t b = 0; b < nbins; ++b)
        ring_thr[b] = (ring_cnt[b] < adaptive_threshold::MIN_RING_PIXELS)
                          ? g_thr
                          : adaptive_threshold::RingThreshold(ring_mean[b], ring_sigma[b], p, z);

    // --- Stage C: the ring threshold, intersected with the classic local-box SNR test ---
    FlagRings(image);

    if (settings.signal_to_noise_threshold <= 0.0f) {
        // No local test asked for: the ring threshold alone decides, as the fixed photon floor
        // alone would in the classic finder.
        output_buffer = ring_bits;
        return;
    }

    // The ring threshold IS the photon floor here, so the local pass must not apply another one.
    SpotFindingSettings local = settings;
    local.photon_count_threshold = 0;
    ImageSpotFinderCPU::Detect(image, local);
    for (size_t i = 0; i < OutputSize(); ++i)
        output_buffer[i] &= ring_bits[i];
}

void AdaptiveSpotFinderCPU::FlagRings(const ImagePreprocessorBuffer &image) {
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const size_t nbins = ring_thr.size();
    const size_t npix = static_cast<size_t>(width) * height;

    std::bitset<32> out = 0;
    for (size_t pxl = 0; pxl < npix; ++pxl) {
        const int32_t v = image[pxl];
        const uint16_t b = pixel_to_bin[pxl];
        bool strong = false;
        if (v == INT32_MAX)
            strong = true;
        else if (v != INT32_MIN && b < nbins && v >= ring_thr[b])
            strong = true;

        const int32_t bit = pxl % 32;
        if (strong)
            out.set(bit);
        if (bit == 31) {
            ring_bits[pxl / 32] = out.to_ulong();
            out.reset();
        }
    }
    if (npix % 32 != 0)
        ring_bits[OutputSize() - 1] = out.to_ulong();
}