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

#include "FrenchWilson.h"

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

#include "../../common/ResolutionShells.h"
#include "gemmi/symmetry.hpp"

namespace {

struct Posterior {
    double mean_I;   // <J>   (posterior mean true intensity)
    double mean_F;   // <|F|> (posterior mean amplitude)
};

// Posterior moments of the true intensity J >= 0 given a measurement I +/- sigma and the Wilson
// prior with mean sigma_wilson. Integrated numerically over J in [0, I + 8 sigma] with a log-shift
// so the exponentials never overflow/underflow. acentric: p(J) ~ exp(-J/S); centric:
// p(J) ~ exp(-J/2S)/sqrt(J).
Posterior integrate_posterior(double I, double sigma, double sigma_wilson, bool centric, int npts) {
    const double inv_2s2 = 1.0 / (2.0 * sigma * sigma);
    const double j_max = std::max(I, 0.0) + 8.0 * sigma;
    const double dj = j_max / npts;

    std::vector<double> logw(npts);
    double max_logw = -std::numeric_limits<double>::infinity();
    for (int i = 0; i < npts; ++i) {
        const double j = (i + 0.5) * dj;
        const double diff = I - j;
        const double log_prior = centric ? (-j / (2.0 * sigma_wilson) - 0.5 * std::log(j))
                                          : (-j / sigma_wilson);
        logw[i] = log_prior - diff * diff * inv_2s2;
        max_logw = std::max(max_logw, logw[i]);
    }

    double sum_w = 0, sum_wI = 0, sum_wF = 0;
    for (int i = 0; i < npts; ++i) {
        const double j = (i + 0.5) * dj;
        const double w = std::exp(logw[i] - max_logw);
        if (!std::isfinite(w))
            continue;
        sum_w += w;
        sum_wI += w * j;
        sum_wF += w * std::sqrt(j);
    }
    if (sum_w <= 0.0) {
        const double j = std::max(I, 0.0);
        return {j, std::sqrt(j)};
    }
    return {sum_wI / sum_w, sum_wF / sum_w};
}

}  // namespace

void ApplyFrenchWilson(std::vector<MergedReflection> &merged, int32_t space_group_number,
                       const FrenchWilsonOptions &opts) {
    auto naive = [](MergedReflection &r) {
        const double ip = std::max(r.I, 0.0f);
        r.F = static_cast<float>(std::sqrt(ip));
        r.sigmaF = (ip > 0.0 && std::isfinite(r.sigma)) ? static_cast<float>(r.sigma / (2.0 * std::sqrt(ip)))
                                                        : NAN;
    };

    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
    if (sg == nullptr || merged.empty()) {
        for (auto &r : merged) naive(r);
        return;
    }
    const gemmi::GroupOps gops = sg->operations();

    float d_min = std::numeric_limits<float>::max(), d_max = 0.0f;
    for (const auto &r : merged)
        if (std::isfinite(r.d) && r.d > 0.0f) {
            d_min = std::min(d_min, r.d);
            d_max = std::max(d_max, r.d);
        }
    if (!(d_min < d_max && d_min > 0.0f)) {
        for (auto &r : merged) naive(r);
        return;
    }

    // Wilson mean intensity <I/epsilon> per resolution shell.
    ResolutionShells shells(d_min * 0.999f, d_max * 1.001f, opts.num_shells);
    std::vector<double> shell_sum(opts.num_shells, 0.0);
    std::vector<int> shell_count(opts.num_shells, 0);
    double global_sum = 0.0;
    int global_count = 0;
    auto epsilon = [&](const MergedReflection &r) {
        return std::max(1, gops.epsilon_factor_without_centering({{r.h, r.k, r.l}}));
    };
    for (const auto &r : merged) {
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f)
            continue;
        const double i_over_eps = r.I / epsilon(r);
        global_sum += i_over_eps;
        ++global_count;
        if (const auto s = shells.GetShell(r.d)) {
            shell_sum[*s] += i_over_eps;
            ++shell_count[*s];
        }
    }
    const double global_mean = global_count > 0 ? std::max(global_sum / global_count, 1e-10) : 1.0;
    std::vector<double> shell_mean(opts.num_shells, global_mean);
    for (int s = 0; s < opts.num_shells; ++s)
        if (shell_count[s] >= opts.min_reflections_per_shell)
            shell_mean[s] = std::max(shell_sum[s] / shell_count[s], 1e-10);

    for (auto &r : merged) {
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f) {
            naive(r);
            continue;
        }
        // Strong reflections: the FW correction is negligible, <|F|> = sqrt(I).
        if (r.I > opts.strong_cutoff * r.sigma) {
            naive(r);
            continue;
        }
        const auto s = shells.GetShell(r.d);
        const double sigma_wilson = epsilon(r) * (s ? shell_mean[*s] : global_mean);
        const bool centric = gops.is_reflection_centric({{r.h, r.k, r.l}});
        const Posterior post = integrate_posterior(r.I, r.sigma, sigma_wilson, centric,
                                                   opts.integration_points);
        r.F = static_cast<float>(post.mean_F);
        r.sigmaF = static_cast<float>(std::sqrt(std::max(0.0, post.mean_I - post.mean_F * post.mean_F)));
    }
}