// 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_MAD(const std::vector<SpotToSave>& xs) {
    std::vector<float> absx;
    absx.reserve(xs.size());
    for (const auto &s: xs) {
        if (s.indexed)
            absx.push_back(std::fabs(s.dist_ewald_sphere));
    }

    if (absx.empty())
        return std::nullopt;

    std::nth_element(absx.begin(), absx.begin() + absx.size() / 2, absx.end());
    float med;
    if (absx.size() % 2 == 1) {
        med = absx[absx.size() / 2];
    } else {
        auto it1 = absx.begin() + (absx.size() / 2 - 1);
        auto it2 = absx.begin() + (absx.size() / 2);
        float a = *it1;
        float b = *it2;
        med = 0.5f * (a + b);
    }
    // Normal consistency factor for MAD
    return 1.4826f * med;
}

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));
}