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

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

#include "CalcISigma.h"
#include "Regression.h"
#include "../../common/ResolutionShells.h"

// Upper bound on a physically plausible macromolecular isotropic B-factor (A^2). A fit that lands
// outside (0, WILSON_B_MAX) comes from a bad frame / bad dataset (too few or mis-indexed reflections,
// a degenerate resolution range) rather than real Debye-Waller falloff, so it is rejected as
// indeterminate rather than reported. Radiation-damaged / low-resolution data rarely exceeds ~120 A^2;
// 200 leaves generous head-room while still excluding the hundreds-of-A^2 garbage.
static constexpr float WILSON_B_MAX = 200.0f;

void CalcISigma(DataMessage &msg) {
    CalcISigma(msg, msg.reflections);
}

void CalcISigma(DataMessage &msg, const std::vector<Reflection> &reflections) {
    if (reflections.empty())
        return;

    const int nshells = 20;
    ResolutionShells shells(1.5, 50.0, nshells);

    std::vector<float> Isigma_sum(nshells);
    std::vector<float> count(nshells);

    for (const auto &r: reflections) {
        auto s = shells.GetShell(r.d);
        if (s && (r.sigma != 0.0)) {
            Isigma_sum[*s] += r.I / r.sigma;
            ++count[*s];
        }
    }

    std::vector<float> result(nshells);
    for (int i = 0; i < nshells; ++i) {
        if (count[i] > 0)
            result[i] = Isigma_sum[i] / count[i];
        else
            result[i] = 0.0f;
    }

    msg.integration_Isigma = result;
    msg.integration_Isigma_one_over_d_square = shells.GetShellMeanOneOverResSq();
}

void CalcWilsonBFactor(DataMessage &msg,
    bool replace_b) {
    CalcWilsonBFactor(msg, msg.reflections, replace_b);
}

void CalcWilsonBFactor(DataMessage &msg,
    const std::vector<Reflection> &reflections,
    bool replace_b) {
    if (reflections.empty())
        return;

    const int nshells = 20;
    ResolutionShells shells(1.5, 6.0, nshells);

    std::vector<float> I_sum(nshells);
    std::vector<float> count(nshells);

    for (const auto& r: reflections) {
        auto s = shells.GetShell(r.d);
        if (s && (r.sigma != 0.0)) {
            I_sum[*s] += r.I;
            ++count[*s];
        }
    }

    std::vector<float> log_I_mean(nshells);

    int32_t valid_shells = nshells;
    for (int i = 0; i < nshells; ++i) {
        if (count[i] > 0 && I_sum[i] > 0) {
            log_I_mean[i] = std::log(I_sum[i] / count[i]);
        } else {
            log_I_mean[i] = 0.0f;
            // First shell that has improper value limits how far the Wilson plot is interpolated
            valid_shells = std::min(valid_shells, i + 1);
        }
    }

    auto shells_mean_one_over_d_square = shells.GetShellMeanOneOverResSq();

    if (replace_b && valid_shells > 2) {
        auto reg_result = regression(shells_mean_one_over_d_square, log_I_mean, valid_shells);
        const float b_est = -2.0f * reg_result.slope;

        // Accept only a well-correlated, physically plausible fit. Occasional bad frames (an indexing
        // glitch, too few reflections) produce a wildly steep Wilson line and a B of several hundred
        // A^2 that pollutes the per-image plot; leaving b_factor unset (rendered as NaN) is better than
        // emitting garbage.
        if (reg_result.r_square > 0.3 && std::isfinite(b_est) && b_est > 0.0f && b_est < WILSON_B_MAX)
             msg.b_factor = b_est;
    }

    msg.integration_B_logI = log_I_mean;
    msg.integration_B_one_over_d_square = shells_mean_one_over_d_square;
}

GlobalWilsonB CalcGlobalWilsonB(const std::vector<MergedReflection> &merged) {
    GlobalWilsonB out;
    // A dataset-wide estimate needs enough reflections to average the shell means; below this the
    // per-image estimate is the only thing on offer and a global number would be meaningless.
    if (merged.size() < 100)
        return out;

    float d_min = std::numeric_limits<float>::infinity(), d_max = 0.0f;
    for (const auto &r : merged) {
        if (std::isfinite(r.I) && 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))
        return out;

    // Below ~4 A the Wilson plot is non-linear (bonding/solvent structure), so when the data extend to
    // lower resolution than that, restrict the fit to d <= 4 A - the standard Wilson-B convention. If
    // the whole dataset is coarser than 4 A, fall back to using all of it.
    constexpr double WILSON_LOW_RES_LIMIT_A = 4.0;
    const float d_low = (d_max > WILSON_LOW_RES_LIMIT_A && d_min < WILSON_LOW_RES_LIMIT_A)
                        ? static_cast<float>(WILSON_LOW_RES_LIMIT_A) : d_max;

    const int nshells = 20;
    ResolutionShells shells(d_min, d_low, nshells);
    std::vector<double> I_sum(nshells, 0.0), sig_sum(nshells, 0.0), count(nshells, 0.0);
    for (const auto &r : merged) {
        if (!std::isfinite(r.I) || r.d <= 0.0f)
            continue;
        auto s = shells.GetShell(r.d);          // reflections coarser than d_low fall outside -> skipped
        if (s) {
            I_sum[*s] += r.I;
            if (std::isfinite(r.sigma) && r.sigma > 0.0f)
                sig_sum[*s] += r.sigma;
            ++count[*s];
        }
    }

    const auto s2 = shells.GetShellMeanOneOverResSq();
    std::vector<float> x, y;
    for (int i = 0; i < nshells; ++i) {
        // Skip empty / net-negative shells, and shells past the signal limit (mean I/sigma < 1). The
        // latter keeps the fit out of the noise floor: without a resolution cut the weakest high-angle
        // shells are background-residual-dominated and flatten the Wilson line, deflating B. This mirrors
        // XDS/ctruncate fitting only over the meaningful range and makes the estimate insensitive to how
        // far the merged data were carried.
        if (count[i] > 0 && I_sum[i] > 0.0 && I_sum[i] > sig_sum[i]) {
            x.push_back(s2[i]);
            y.push_back(std::log(static_cast<float>(I_sum[i] / count[i])));
        }
    }
    if (x.size() < 3)
        return out;

    const auto reg = regression(x, y, x.size());
    const double b = -2.0 * reg.slope;          // <I> ~ exp(-2 B s^2), s^2 = 1/(4 d^2), x = 1/d^2
    out.n_shells = static_cast<int>(x.size());
    out.correlation = std::sqrt(std::clamp(static_cast<double>(reg.r_square), 0.0, 1.0));
    if (std::isfinite(b) && b > 0.0 && b < WILSON_B_MAX)
        out.b = b;
    return out;
}