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

#include "FitProfileRadius.h"
#include <algorithm>  // std::nth_element
#include <cmath>      // std::fabs

std::optional<float> FitProfileRadius(const std::vector<SpotToSave>& spots,
                                      float bandwidth_sigma, float wavelength_A) {
    double sum_squares = 0.0;     // measured excitation-error variance (sum dist_ewald^2)
    double sum_bw_var = 0.0;      // energy-bandwidth contribution to subtract out
    int count = 0;

    // Fitted from the strongest spots only. This is an RMS over whatever spots were kept, and weaker
    // spots sit further off the Ewald sphere, so it grows with the depth of the list: measured over a
    // 150 -> unlimited spot budget it rises ~60% on both a hard and a clean rotation dataset. That
    // made it a function of the INDEXING budget (--max-spots) rather than of the crystal. Its one
    // consumer treats it as a membership gate (ewald_dist_cutoff = 2x it), where reflections at the
    // cutoff carry ~zero partiality, so the integrated data barely moved (0.2% of partials across
    // that range) - but it is also reported per image as a diagnostic, where a number that slides
    // with an unrelated setting is simply misleading. FilterSpotsByCount leaves the list
    // strongest-first, so taking the head selects the spots a smaller --max-spots would.
    constexpr size_t PROFILE_RADIUS_FIT_SPOTS = 250;
    const size_t n_fit = std::min(spots.size(), PROFILE_RADIUS_FIT_SPOTS);

    for (size_t si = 0; si < n_fit; ++si) {
        const auto &s = spots[si];
        if (!s.indexed)
            continue;
        sum_squares += static_cast<double>(s.dist_ewald_sphere) * s.dist_ewald_sphere;
        // The energy bandwidth smears each reflection radially by sigma_bw = bandwidth_sigma*|recip_z|
        // = bandwidth_sigma*lambda/(2 d^2) (the same term prediction re-adds per reflection, ~1/d^2 so
        // largest at high resolution). Deconvolve it from the measured spread so the profile radius is
        // the *intrinsic* mosaicity+divergence width and bandwidth is not double-counted at prediction.
        if (bandwidth_sigma > 0.0f && s.d_A > 0.0f) {
            const double sigma_bw = bandwidth_sigma * wavelength_A / (2.0 * static_cast<double>(s.d_A) * s.d_A);
            sum_bw_var += sigma_bw * sigma_bw;
        }
        count++;
    }

    if (count == 0)
        return std::nullopt;

    const double variance = std::max(0.0, (sum_squares - sum_bw_var) / count);
    return static_cast<float>(std::sqrt(variance));
}