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

#include "ProfileIntegrate2D.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <vector>

#include "../../common/JFJochException.h"

namespace {

constexpr int N_SHELL = 6;
constexpr double STRONG_I_OVER_SIGMA = 5.0;
constexpr int MIN_STRONG_PER_SHELL = 30;

// Rough box-sum of one reflection: total intensity over the inner disk (r1) minus the mean of the
// background ring (r2..r3), like BraggIntegrate2D. Used to seed the fit, pick strong spots for
// profile learning, and provide the per-reflection background. ok=false when the inner disk is
// clipped by a gap/edge/bad pixel (same rejection as the box-sum integrator).
struct Rough {
    double I = 0.0, sigma = NAN, bkg = 0.0;
    int cx = 0, cy = 0, shell = -1;
    bool ok = false, strong = false;
};

template<class T>
Rough BoxSum(const T *image, size_t xpixel, size_t ypixel, int64_t special, int64_t saturation,
             const Reflection &r, float r1_sq, float r2_sq, float r3_sq, float r3, float min_sigma_ratio,
             bool apply_bkg_clip) {
    Rough out;
    const int64_t x0 = std::max<int64_t>(0, std::floor(r.predicted_x - r3 - 1.0));
    const int64_t x1 = std::min<int64_t>(xpixel - 1, std::ceil(r.predicted_x + r3 + 1.0));
    const int64_t y0 = std::max<int64_t>(0, std::floor(r.predicted_y - r3 - 1.0));
    const int64_t y1 = std::min<int64_t>(ypixel - 1, std::ceil(r.predicted_y + r3 + 1.0));

    int64_t I_sum = 0, n_inner = 0, n_inner_valid = 0;
    double bkg_sum = 0.0;
    int n_bkg = 0;
    std::vector<double> bkg_vals;
    bkg_vals.reserve(128);
    for (int64_t y = y0; y <= y1; ++y)
        for (int64_t x = x0; x <= x1; ++x) {
            const double d2 = (x - r.predicted_x) * (x - r.predicted_x) + (y - r.predicted_y) * (y - r.predicted_y);
            const auto px = image[y * xpixel + x];
            if (d2 < r1_sq) {
                ++n_inner;
                if (px == special || px == special + 1 || px == saturation || px == saturation - 1) continue;
                I_sum += px;
                ++n_inner_valid;
            } else if (d2 >= r2_sq && d2 < r3_sq) {
                if (px == special || px == special + 1 || px == saturation || px == saturation - 1) continue;
                bkg_sum += static_cast<double>(px);
                ++n_bkg;
                bkg_vals.push_back(static_cast<double>(px));
            }
        }
    if (n_inner_valid != n_inner || n_bkg <= 5)
        return out;
    out.bkg = bkg_sum / n_bkg;
    // One high-outlier sigma-clip pass on the background ring. A bandwidth-streaked high-resolution spot
    // (or a neighbour) leaks into the annulus and biases the mean high, which over-subtracts and drives
    // the merged high-resolution intensities negative; rejecting pixels above mean + 3*sqrt(mean) removes
    // that contamination (and barely touches a clean Poisson background, where ~0.1% exceed the cut).
    // Stills-only: on rotation (no bandwidth streak) it clips legitimate high background pixels and biases
    // the mean low, slightly hurting weak (anomalous) intensities, so it is gated off there.
    if (apply_bkg_clip) {
        const double thr = out.bkg + 3.0 * std::sqrt(std::max(out.bkg, 1.0));
        double s = 0.0;
        int n = 0;
        for (const double v : bkg_vals) if (v <= thr) { s += v; ++n; }
        if (n > 5) out.bkg = s / n;
    }
    out.I = static_cast<double>(I_sum) - static_cast<double>(n_inner) * out.bkg;
    out.sigma = std::max(1.0, out.I * min_sigma_ratio);
    if (I_sum > 0)
        out.sigma = std::max(out.sigma, std::sqrt(static_cast<double>(I_sum)));
    out.cx = static_cast<int>(std::lround(r.predicted_x));
    out.cy = static_cast<int>(std::lround(r.predicted_y));
    out.ok = true;
    out.strong = out.sigma > 0.0 && out.I / out.sigma >= STRONG_I_OVER_SIGMA;
    return out;
}

template<class T>
std::vector<Reflection> ProfileIntegrateInternal(const DiffractionExperiment &experiment,
                                                 const CompressedImage &image,
                                                 const std::vector<Reflection> &predicted, size_t npredicted,
                                                 int64_t special, int64_t saturation, int64_t image_number) {
    const auto settings = experiment.GetBraggIntegrationSettings();
    const auto geom = experiment.GetDiffractionGeometry();
    const bool empirical = settings.GetIntegrator() == IntegratorMode::ProfileEmpirical;

    const float r1_sq = settings.GetR1() * settings.GetR1();
    const float r2 = settings.GetR2(), r2_sq = r2 * r2;
    const float r3 = settings.GetR3(), r3_sq = r3 * r3;
    const float min_sigma_ratio = settings.GetMinimumSigmaInRegardsToI();
    // The background-ring sigma-clip de-biases bandwidth-streaked stills but can clip legitimate
    // background on rotation (no streaks) and bias it low; gate it to stills (a bandwidth is set).
    const bool apply_bkg_clip = experiment.GetBandwidthFWHM().value_or(0.0f) > 0.0f;
    const int R = static_cast<int>(std::ceil(r2));   // profile grid half-size
    const int G = 2 * R + 1, GG = G * G;
    auto idx = [G, R](int dx, int dy) { return (dy + R) * G + (dx + R); };

    std::vector<uint8_t> buffer;
    const auto *ptr = reinterpret_cast<const T *>(image.GetUncompressedPtr(buffer));
    const size_t xpixel = image.GetWidth(), ypixel = image.GetHeight();

    // --- Pass A: box-sum every reflection (rough I, background, centre, strong flag). ---
    std::vector<Rough> rough(npredicted);
    double inv_d2_min = std::numeric_limits<double>::max(), inv_d2_max = 0.0;
    for (size_t i = 0; i < npredicted; ++i) {
        rough[i] = BoxSum<T>(ptr, xpixel, ypixel, special, saturation, predicted[i],
                             r1_sq, r2_sq, r3_sq, r3, min_sigma_ratio, apply_bkg_clip);
        if (rough[i].ok && predicted[i].d > 0.0f) {
            const double inv_d2 = 1.0 / (predicted[i].d * predicted[i].d);
            inv_d2_min = std::min(inv_d2_min, inv_d2);
            inv_d2_max = std::max(inv_d2_max, inv_d2);
        }
    }
    auto shell_of = [&](float d) {
        if (!(d > 0.0f) || inv_d2_max <= inv_d2_min) return 0;
        const double t = (1.0 / (d * d) - inv_d2_min) / (inv_d2_max - inv_d2_min);
        return std::clamp(static_cast<int>(t * N_SHELL), 0, N_SHELL - 1);
    };
    for (size_t i = 0; i < npredicted; ++i)
        if (rough[i].ok)
            rough[i].shell = shell_of(predicted[i].d);

    // --- Learn the profile per shell from the strong spots (+ a global fallback). Each strong
    //     spot contributes its background-subtracted, intensity-normalised, centroid-aligned grid. ---
    std::vector<std::vector<double>> shell_grid(N_SHELL, std::vector<double>(GG, 0.0));
    std::vector<int> shell_n(N_SHELL, 0);
    std::vector<double> global_grid(GG, 0.0);
    int global_n = 0;
    for (size_t i = 0; i < npredicted; ++i) {
        const auto &rh = rough[i];
        if (!rh.ok || !rh.strong || rh.I <= 0.0) continue;
        for (int dy = -R; dy <= R; ++dy)
            for (int dx = -R; dx <= R; ++dx) {
                const int64_t x = rh.cx + dx, y = rh.cy + dy;
                if (x < 0 || y < 0 || x >= static_cast<int64_t>(xpixel) || y >= static_cast<int64_t>(ypixel)) continue;
                const auto px = ptr[y * xpixel + x];
                if (px == special || px == special + 1 || px == saturation || px == saturation - 1) continue;
                const double v = (static_cast<double>(px) - rh.bkg) / rh.I;
                shell_grid[rh.shell][idx(dx, dy)] += v;
                global_grid[idx(dx, dy)] += v;
            }
        ++shell_n[rh.shell];
        ++global_n;
    }

    // Build a normalised profile (sum = 1): the empirical average, or an isotropic Gaussian of the
    // same (measured) second moment - the spots are ~round in the detector plane, so radial/tangential
    // anisotropy was measured and found to add nothing (the elongation is in the discarded rocking
    // direction), and a single width is kept for simplicity.
    auto build_profile = [&](const std::vector<double> &grid, int n) {
        std::vector<double> P(GG, 0.0);
        if (n <= 0) return P;
        double sum = 0.0, m2 = 0.0, m2w = 0.0;
        for (int dy = -R; dy <= R; ++dy)
            for (int dx = -R; dx <= R; ++dx) {
                const double g = std::max(0.0, grid[idx(dx, dy)]);
                m2 += g * (dx * dx + dy * dy);
                m2w += g;
                sum += g;
                if (empirical) P[idx(dx, dy)] = g;
            }
        if (sum <= 0.0) return P;
        if (empirical) {
            for (double &p : P) p /= sum;
        } else {
            const double sigma2 = std::max(0.25, (m2w > 0.0 ? m2 / m2w : 1.0) / 2.0);   // <r^2> = 2 sigma^2 (2D)
            double gsum = 0.0;
            for (int dy = -R; dy <= R; ++dy)
                for (int dx = -R; dx <= R; ++dx) {
                    const double g = std::exp(-(dx * dx + dy * dy) / (2.0 * sigma2));
                    P[idx(dx, dy)] = g;
                    gsum += g;
                }
            for (double &p : P) p /= gsum;
        }
        return P;
    };
    const std::vector<double> global_P = build_profile(global_grid, global_n);
    std::vector<std::vector<double>> shell_P(N_SHELL);
    for (int s = 0; s < N_SHELL; ++s)
        shell_P[s] = shell_n[s] >= MIN_STRONG_PER_SHELL ? build_profile(shell_grid[s], shell_n[s]) : global_P;

    // Isotropic (intrinsic) width per shell, used to seed the radially-elongated profile below.
    auto measure_sigma2 = [&](const std::vector<double> &grid) {
        double m2 = 0.0, m2w = 0.0;
        for (int dy = -R; dy <= R; ++dy)
            for (int dx = -R; dx <= R; ++dx) {
                const double g = std::max(0.0, grid[idx(dx, dy)]);
                m2 += g * (dx * dx + dy * dy);
                m2w += g;
            }
        return m2w > 0.0 ? std::max(0.25, (m2 / m2w) / 2.0) : 1.0;
    };
    const double global_sigma2 = global_n > 0 ? measure_sigma2(global_grid) : 1.0;
    std::vector<double> shell_sigma2(N_SHELL, global_sigma2);
    for (int s = 0; s < N_SHELL; ++s)
        if (shell_n[s] >= MIN_STRONG_PER_SHELL)
            shell_sigma2[s] = measure_sigma2(shell_grid[s]);

    // The energy bandwidth smears a reflection RADIALLY by sigma_bw = bandwidth_sigma * R_px (R_px =
    // distance from the beam centre, so large at high resolution). The isotropic profile then both
    // mis-weights it and (on the fixed grid) clips its radial tail. In ellipse mode, fit each
    // reflection with a per-reflection Gaussian elongated only along the radial direction
    // (sigma^2_radial = sigma^2_intrinsic + sigma_bw^2, sigma^2_tangential = sigma^2_intrinsic) on a
    // grid grown to hold the streak - no extra tangential background, unlike an isotropic widening.
    const double bw_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f;
    const float beam_x = geom.GetBeamX_pxl(), beam_y = geom.GetBeamY_pxl();
    const bool use_ellipse = !empirical && bw_sigma > 0.0;   // only acts when a bandwidth is set

    // --- Pass B: profile-fit each reflection (Kabsch, de-biased variance v = B + I*P; iterate). ---
    std::vector<Reflection> out;
    out.reserve(npredicted);
    const auto pol = experiment.GetPolarizationFactor();
    std::vector<double> Pbuf;
    for (size_t i = 0; i < npredicted; ++i) {
        const auto &rh = rough[i];
        if (!rh.ok) continue;
        const int sh = rh.shell < 0 ? 0 : rh.shell;

        int Rf = R;
        const std::vector<double> *Pvec = &shell_P[sh];
        if (use_ellipse) {
            const double rx = predicted[i].predicted_x - beam_x, ry = predicted[i].predicted_y - beam_y;
            const double Rpx = std::hypot(rx, ry), sbw = bw_sigma * Rpx;
            // Only elongate where the streak is genuinely larger than the intrinsic spot (high
            // resolution); at low/mid resolution the smear is sub-pixel and elongating just adds noise.
            if (Rpx > 1e-6 && sbw * sbw > shell_sigma2[sh]) {
                const double ux = rx / Rpx, uy = ry / Rpx;
                const double s2r = shell_sigma2[sh] + sbw * sbw, s2t = shell_sigma2[sh];
                Rf = std::min(3 * R, static_cast<int>(std::ceil(r2 + 2.0 * std::sqrt(s2r))));
                const int Gf = 2 * Rf + 1;
                Pbuf.assign(static_cast<size_t>(Gf) * Gf, 0.0);
                double gs = 0.0;
                for (int dy = -Rf; dy <= Rf; ++dy)
                    for (int dx = -Rf; dx <= Rf; ++dx) {
                        const double rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
                        const double g = std::exp(-rad * rad / (2.0 * s2r) - tn * tn / (2.0 * s2t));
                        Pbuf[(dy + Rf) * Gf + (dx + Rf)] = g;
                        gs += g;
                    }
                for (double &p : Pbuf) p /= gs;
                Pvec = &Pbuf;
            }
        }

        const int Gf = 2 * Rf + 1;
        const double B = std::max(rh.bkg, 1.0);
        double I = rh.I, den = 0.0;
        for (int iter = 0; iter < 4; ++iter) {
            double num = 0.0;
            den = 0.0;
            for (int dy = -Rf; dy <= Rf; ++dy)
                for (int dx = -Rf; dx <= Rf; ++dx) {
                    const double Pp = (*Pvec)[(dy + Rf) * Gf + (dx + Rf)];
                    if (Pp <= 0.0) continue;
                    const int64_t x = rh.cx + dx, y = rh.cy + dy;
                    if (x < 0 || y < 0 || x >= static_cast<int64_t>(xpixel) || y >= static_cast<int64_t>(ypixel)) continue;
                    const auto px = ptr[y * xpixel + x];
                    if (px == special || px == special + 1 || px == saturation || px == saturation - 1) continue;
                    const double v = B + std::max(0.0, I) * Pp;
                    num += Pp * (static_cast<double>(px) - rh.bkg) / v;
                    den += Pp * Pp / v;
                }
            if (den > 0.0) I = num / den; else break;
        }
        if (!(den > 0.0)) continue;

        Reflection refl = predicted[i];
        refl.I = static_cast<float>(I);
        refl.sigma = static_cast<float>(std::sqrt(1.0 / den));
        refl.bkg = static_cast<float>(rh.bkg);
        refl.observed = true;
        if (pol)
            refl.rlp /= geom.CalcAzIntPolarizationCorr(refl.predicted_x, refl.predicted_y, pol.value());
        refl.image_scale_corr = refl.rlp / refl.partiality;
        refl.image_number = static_cast<float>(image_number);
        out.push_back(refl);
    }
    return out;
}

}   // namespace

std::vector<Reflection> ProfileIntegrate2D(const DiffractionExperiment &experiment,
                                           const CompressedImage &image,
                                           const std::vector<Reflection> &predicted, size_t npredicted,
                                           int64_t image_number) {
    if (image.GetCompressedSize() == 0 || predicted.empty())
        return {};
    switch (image.GetMode()) {
        case CompressedImageMode::Int8:
            return ProfileIntegrateInternal<int8_t>(experiment, image, predicted, npredicted, INT8_MIN, INT8_MAX, image_number);
        case CompressedImageMode::Int16:
            return ProfileIntegrateInternal<int16_t>(experiment, image, predicted, npredicted, INT16_MIN, INT16_MAX, image_number);
        case CompressedImageMode::Int32:
            return ProfileIntegrateInternal<int32_t>(experiment, image, predicted, npredicted, INT32_MIN, INT32_MAX, image_number);
        case CompressedImageMode::Uint8:
            return ProfileIntegrateInternal<uint8_t>(experiment, image, predicted, npredicted, UINT8_MAX, UINT8_MAX, image_number);
        case CompressedImageMode::Uint16:
            return ProfileIntegrateInternal<uint16_t>(experiment, image, predicted, npredicted, UINT16_MAX, UINT16_MAX, image_number);
        case CompressedImageMode::Uint32:
            return ProfileIntegrateInternal<uint32_t>(experiment, image, predicted, npredicted, UINT32_MAX, UINT32_MAX, image_number);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image mode not supported");
    }
}