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

    for (const auto &s: spots) {
        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));
}