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

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

void CalcISigma(DataMessage &msg) {
    if (msg.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: msg.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) {
    if (msg.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: msg.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);

        if (reg_result.r_square > 0.3)
             msg.b_factor = -2.0f * reg_result.slope;
    }

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