// 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 <unordered_map>

#include "AdaptiveSpotFinderCPU.h"

namespace {

// Number of background pixels a ring needs before its own statistics are trusted; sparser rings
// (detector corners, heavily masked, innermost) fall back to the whole-frame background.
constexpr int64_t MIN_RING_PIXELS = 40;

// Inverse standard-normal CDF (Acklam's rational approximation, ~1e-9 accuracy). Only called once
// per frame, so accuracy over speed.
double NormalQuantile(double p) {
    if (p <= 0.0) return -40.0;
    if (p >= 1.0) return 40.0;
    static const double a[] = {-3.969683028665376e+01, 2.209460984245205e+02, -2.759285104469687e+02,
                                1.383577518672690e+02, -3.066479806614716e+01, 2.506628277459239e+00};
    static const double b[] = {-5.447609879822406e+01, 1.615858368580409e+02, -1.556989798598866e+02,
                                6.680131188771972e+01, -1.328068155288572e+01};
    static const double c[] = {-7.784894002430293e-03, -3.223964580411365e-01, -2.400758277161838e+00,
                               -2.549732539343734e+00, 4.374664141464968e+00, 2.938163982698783e+00};
    static const double d[] = {7.784695709041462e-03, 3.224671290700398e-01, 2.445134137142996e+00,
                               3.754408661907416e+00};
    const double plow = 0.02425, phigh = 1.0 - 0.02425;
    if (p < plow) {
        double q = std::sqrt(-2.0 * std::log(p));
        return (((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) /
               ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1.0);
    } else if (p <= phigh) {
        double q = p - 0.5, r = q*q;
        return (((((a[0]*r+a[1])*r+a[2])*r+a[3])*r+a[4])*r+a[5])*q /
               (((((b[0]*r+b[1])*r+b[2])*r+b[3])*r+b[4])*r+1.0);
    } else {
        double q = std::sqrt(-2.0 * std::log(1.0 - p));
        return -(((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) /
                ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1.0);
    }
}

// Smallest integer count whose Poisson(mu) upper tail P(X >= k) <= p. This is the correct
// significance floor while the background is countable (it carries the sqrt(mu) shot-noise
// implicitly, so a bright low-resolution ring gets a high threshold). It DEGENERATES at mu -> 0
// (a single photon on a zero background is "significant"), which is why it is max'd with a
// read-noise-floored Gaussian arm by the caller. Short-circuits to Gaussian for large mu.
float PoissonThreshold(double mu, double p, double z) {
    if (mu > 50.0)
        return static_cast<float>(mu + z * std::sqrt(mu));
    if (mu < 1e-6) mu = 1e-6;
    const double target = 1.0 - p;
    double pmf = std::exp(-mu);
    double cdf = pmf;
    int k = 0;
    while (cdf < target && k < 1000) {
        ++k;
        pmf *= mu / k;
        cdf += pmf;
    }
    return static_cast<float>(k + 1);
}

} // namespace

AdaptiveSpotFinderCPU::AdaptiveSpotFinderCPU(const AzimuthalIntegrationMapping &in_mapping)
    : ImageSpotFinder(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.0);
    ring_sum2.assign(nbins, 0.0);
    ring_cnt.assign(nbins, 0);
    ring_mean.assign(nbins, 0.0f);
    ring_sigma.assign(nbins, 0.0f);
    ring_thr.assign(nbins, 0.0f);
}

// 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.0);
    std::fill(ring_sum2.begin(), ring_sum2.end(), 0.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<double>(v) * v;
        ring_cnt[b] += 1;
    }

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

std::vector<DiffractionSpot> AdaptiveSpotFinderCPU::Run(const ImagePreprocessorBuffer &image,
                                                        const SpotFindingSettings &settings,
                                                        const std::vector<bool> &res_mask) {
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const size_t nbins = ring_sum.size();
    const size_t npix = static_cast<size_t>(width) * height;

    // --- 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 += ring_sum[b];
        g_sum2 += ring_sum2[b];
    }
    if (n_total == 0)
        return {};

    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>(NormalQuantile(1.0 - p));

    // A ring's threshold is background mean + z sigmas. sigma combines the ring's own (peak-excluded)
    // scatter with an excess-noise floor READ: near-zero-background rings scatter MORE than pure
    // Poisson (charge sharing / read noise / occasional spurious low counts), so a per-ring sigma
    // alone collapses toward zero on empty high-resolution rings and the threshold would flood. READ
    // is a detector-level photon-scale constant (the same for every dataset -- it is NOT the
    // per-dataset knob), so the operating point still self-calibrates through mean and sigma while
    // staying physical where the background vanishes.
    const float READ = 1.0f;
    auto ring_threshold = [&](float mean, float sigma) {
        // Poisson significance (correct where the background is countable) floored by a
        // read-noise-aware Gaussian arm (which alone survives mean -> 0, where Poisson degenerates
        // to "one photon is significant" and would flood the empty high-resolution rings).
        const float gauss = mean + z * std::sqrt(sigma * sigma + READ * READ);
        const float poisson = PoissonThreshold(mean, static_cast<double>(p), static_cast<double>(z));
        return std::max(gauss, poisson);
    };

    // 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 = ring_threshold(static_cast<float>(g_mean), static_cast<float>(g_sigma));

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

    // --- Stage C: flag strong pixels into the bit buffer (value >= ring threshold) ---
    for (size_t i = 0; i < OutputSize(); ++i)
        output_buffer[i] = 0;

    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) {
            output_buffer[pxl / 32] = out.to_ulong();
            out.reset();
        }
    }
    if (npix % 32 != 0)
        output_buffer[OutputSize() - 1] = out.to_ulong();

    // --- Stage D: connected components + resolution mask + min/max-pix (shared with classic path) ---
    return ExtractSpots(image, settings, res_mask);
}