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

#include "BeamCenterFFT.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>

#include "BeamCenterFFTCPU.h"
#include "BeamCenterFFTEngine.h"
#include "../../common/CUDAWrapper.h"
#ifdef JFJOCH_USE_CUDA
#include "BeamCenterFFTGPU.h"
#include "../../common/JFJochException.h"
#endif

namespace {

// The image prepared for scoring: hot pixels clipped at the 99.9th percentile, the median
// background level removed, negatives clamped, masked pixels zero. The final subtraction of the
// valid-pixel mean changes nothing mathematically - the masked Pearson is exactly invariant under
// a global shift of the image (both D*C - S^2 and D*Q - S^2 cancel the shift) - but it removes
// most of the large-term cancellation that single-precision FFTs would otherwise have to survive.
struct Prepared {
    std::vector<float> a;   // image, 0 on masked pixels
    std::vector<float> m;   // 1 on valid pixels, 0 on masked
    double variance = 0.0;  // of a over the valid pixels, after the whole preparation
};

Prepared PrepareImage(const std::vector<float> &mean) {
    Prepared out;
    const size_t n = mean.size();
    out.a.assign(n, 0.0f);
    out.m.assign(n, 0.0f);

    std::vector<float> valid;
    valid.reserve(n);
    for (float v : mean)
        if (std::isfinite(v))
            valid.push_back(v);
    if (valid.size() < 2)
        return out;

    auto kth = [&valid](size_t k) {
        std::nth_element(valid.begin(), valid.begin() + k, valid.end());
        return valid[k];
    };
    const float clip_hi = kth(static_cast<size_t>(0.999 * (valid.size() - 1)));
    const float median = kth(valid.size() / 2);

    double sum = 0.0;
    size_t count = 0;
    for (size_t i = 0; i < n; i++) {
        if (!std::isfinite(mean[i]))
            continue;
        const float v = std::max(0.0f, std::min(std::max(mean[i], 0.0f), clip_hi) - median);
        out.a[i] = v;
        out.m[i] = 1.0f;
        sum += v;
        count++;
    }
    const float shift = static_cast<float>(sum / static_cast<double>(count));
    double var = 0.0;
    for (size_t i = 0; i < n; i++)
        if (out.m[i] != 0.0f) {
            out.a[i] -= shift;
            var += static_cast<double>(out.a[i]) * static_cast<double>(out.a[i]);
        }
    out.variance = var / static_cast<double>(count);
    return out;
}

std::vector<BeamCenterFFTCandidate> Shortlist1D(std::vector<float> &surface, float nms_pxl,
                                                int budget, bool is_x) {
    std::vector<BeamCenterFFTCandidate> out;
    const int64_t n = static_cast<int64_t>(surface.size());
    const int64_t d = std::lround(2.0f * nms_pxl);
    for (int k = 0; k < budget; k++) {
        int64_t best = -1;
        float best_v = -std::numeric_limits<float>::infinity();
        for (int64_t i = 0; i < n; i++)
            if (surface[i] > best_v) {
                best_v = surface[i];
                best = i;
            }
        if (best < 0 || !std::isfinite(best_v))
            break;
        const float c = static_cast<float>(best) / 2.0f;
        const float nan = std::numeric_limits<float>::quiet_NaN();
        out.push_back(is_x ? BeamCenterFFTCandidate{c, nan, best_v}
                           : BeamCenterFFTCandidate{nan, c, best_v});
        for (int64_t i = std::max<int64_t>(0, best - d); i <= std::min(n - 1, best + d); i++)
            surface[i] = -std::numeric_limits<float>::infinity();
    }
    return out;
}

float Margin(const std::vector<BeamCenterFFTCandidate> &peaks) {
    if (peaks.size() < 2 || peaks[0].score <= 0.0f)
        return 0.0f;
    return (peaks[0].score - peaks[1].score) / peaks[0].score;
}

// The same score for a 1D line mirror, whose accumulators were summed over the other coordinate.
std::vector<float> PearsonLine(const BeamCenterConvSurfaces1D &conv, float min_pair_fraction,
                               double global_variance) {
    const size_t n = conv.C.size();
    double d_max = 0.0;
    for (double d : conv.D)
        d_max = std::max(d_max, d);
    const double lim = min_pair_fraction * d_max;
    std::vector<float> r(n, -std::numeric_limits<float>::infinity());
    for (size_t i = 0; i < n; i++) {
        const double num = conv.D[i] * conv.C[i] - conv.S[i] * conv.S[i];
        const double den = conv.D[i] * conv.Q[i] - conv.S[i] * conv.S[i];
        // Same no-evidence floor as the 2D surface: see BEAM_CENTER_VARIANCE_FLOOR.
        if (conv.D[i] > lim
            && den > conv.D[i] * conv.D[i] * BEAM_CENTER_VARIANCE_FLOOR * global_variance)
            r[i] = static_cast<float>(num / den);
    }
    return r;
}

} // namespace

std::vector<float> BeamCenterPointScore(const BeamCenterConvSurfaces2D &conv,
                                        float min_pair_fraction, double global_variance) {
    const size_t n = conv.C.size();
    float d_max = 0.0f;
    for (size_t i = 0; i < n; i++)
        d_max = std::max(d_max, conv.D[i]);
    const double lim = min_pair_fraction * d_max;
    std::vector<float> r(n, -std::numeric_limits<float>::infinity());
    for (size_t i = 0; i < n; i++) {
        const double d = conv.D[i];
        if (d <= lim)
            continue;
        const double s = conv.S[i];
        const double num = d * conv.C[i] - s * s;
        const double den = d * conv.Q[i] - s * s;
        if (den > d * d * BEAM_CENTER_VARIANCE_FLOOR * global_variance)
            r[i] = static_cast<float>(num / den);
    }
    return r;
}

std::vector<BeamCenterFFTCandidate> BeamCenterShortlist2D(std::vector<float> &surface, int64_t h,
                                                          int64_t w, float nms_pxl, int budget) {
    std::vector<BeamCenterFFTCandidate> out;
    const int64_t d = std::lround(2.0f * nms_pxl);
    for (int k = 0; k < budget; k++) {
        int64_t best = -1;
        float best_v = -std::numeric_limits<float>::infinity();
        for (int64_t i = 0; i < h * w; i++)
            if (surface[i] > best_v) {
                best_v = surface[i];
                best = i;
            }
        if (best < 0 || !std::isfinite(best_v))
            break;
        const int64_t iy = best / w, ix = best % w;
        out.push_back({static_cast<float>(ix) / 2.0f, static_cast<float>(iy) / 2.0f, best_v});
        for (int64_t y = std::max<int64_t>(0, iy - d); y <= std::min(h - 1, iy + d); y++)
            for (int64_t x = std::max<int64_t>(0, ix - d); x <= std::min(w - 1, ix + d); x++)
                surface[static_cast<size_t>(y * w + x)] =
                        -std::numeric_limits<float>::infinity();
    }
    return out;
}

// The default route, and the whole of the CPU path: the four convolutions come back, and the
// combination and the search run here.
std::vector<BeamCenterFFTCandidate>
BeamCenterFFTEngine::PointShortlist(const std::vector<float> &a, const std::vector<float> &m,
                                    int64_t h, int64_t w, const BeamCenterFFTSettings &settings,
                                    double global_variance) {
    const auto conv = PointSurfaces(a, m, h, w);
    auto surface = BeamCenterPointScore(conv, settings.min_pair_fraction, global_variance);
    return BeamCenterShortlist2D(surface, 2 * h, 2 * w, settings.nms_radius_pxl,
                                 settings.candidates_point);
}

int64_t BeamCenterFFTPadSize(int64_t n) {
    for (int64_t k = n;; ++k) {
        int64_t v = k;
        for (int p : {2, 3, 5, 7})
            while (v % p == 0)
                v /= p;
        if (v == 1)
            return k;
    }
}

BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height,
                                       const std::vector<float> &mean,
                                       const BeamCenterFFTSettings &settings,
                                       BeamCenterFFTEngine &engine) {
    if (width <= 0 || height <= 0
        || mean.size() != static_cast<size_t>(width) * static_cast<size_t>(height))
        throw std::runtime_error("BeamCenterFFT: image dimensions do not match the projection");

    BeamCenterFFTResult result;
    const Prepared prep = PrepareImage(mean);

    result.point = engine.PointShortlist(prep.a, prep.m, height, width, settings, prep.variance);

    {
        const auto conv_y = engine.LineSurfaces(prep.a, prep.m, height, width,
                                                BeamCenterMirror::Rows);
        auto ry = PearsonLine(conv_y, settings.min_pair_fraction, prep.variance);
        result.line_y = Shortlist1D(ry, settings.nms_radius_pxl, settings.candidates_line, false);

        const auto conv_x = engine.LineSurfaces(prep.a, prep.m, height, width,
                                                BeamCenterMirror::Columns);
        auto rx = PearsonLine(conv_x, settings.min_pair_fraction, prep.variance);
        result.line_x = Shortlist1D(rx, settings.nms_radius_pxl, settings.candidates_line, true);
    }

    result.margin_point = Margin(result.point);
    result.margin_line_x = Margin(result.line_x);
    result.margin_line_y = Margin(result.line_y);
    return result;
}

BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height,
                                       const std::vector<float> &mean,
                                       const BeamCenterFFTSettings &settings) {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0 && BeamCenterFFTGPU::FitsInDeviceMemory(width, height)) {
        BeamCenterFFTGPU gpu;
        try {
            return BeamCenterFFTScore(width, height, mean, settings, gpu);
        } catch (const JFJochException &) {
            // The card is shared with this run's analysis workers, so the memory the check above
            // saw can be gone by the time it is asked for. A capture is not worth failing a run
            // over: score it on the CPU instead.
        }
    }
#endif
    BeamCenterFFTCPU cpu;
    return BeamCenterFFTScore(width, height, mean, settings, cpu);
}

BeamCenterFFTResult FindBeamCenterFFT(const DiffractionExperiment &experiment,
                                      const std::vector<float> &mean,
                                      const BeamCenterFFTSettings &settings) {
    // The pre-scan projection is in converted geometry (ShadowFinder's frame).
    return BeamCenterFFTScore(experiment.GetXPixelsNumConv(), experiment.GetYPixelsNumConv(),
                              mean, settings);
}