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

#include "BraggIntegrationStats.h"
#include "Regression.h"
#include <cmath>

BraggIntegrationStats::BraggIntegrationStats(float d_max_A, float d_min_A, int32_t in_nshells) {
    d_min = d_min_A;
    d_max = d_max_A;
    one_over_dmax = 1.0f / (d_max);
    one_over_dmin = 1.0f / (d_min);
    nshells = in_nshells;
    I.resize(nshells);
    I_sigma.resize(nshells);
    count.resize(nshells);
}

void BraggIntegrationStats::AddReflection(const Reflection &r) {
    if (r.d == 0.0f)
        return;

    float one_over_d = 1.0f / (r.d);
    int64_t shell = std::lround((one_over_d - one_over_dmax) / (one_over_dmin - one_over_dmax)
            * static_cast<float>(nshells));
    if (shell >= 0 && shell < nshells) {
        I.at(shell) += r.I;
        I_sigma.at(shell) += r.I / r.sigma;
        count.at(shell) += 1;
    }
}

std::optional<float> BraggIntegrationStats::BFactor() {
    wilson_plot.clear();
    wilson_plot_legend.clear();

    for (int i = 0; i < nshells; i++) {
        float one_over_d = one_over_dmax
                           + (static_cast<float>(i) + 0.5) * (one_over_dmin - one_over_dmax) / static_cast<float>(nshells);
        wilson_plot_legend.push_back(one_over_d);
        if (count[i] == 0 )
            plot_Isigma.push_back(NAN);
        else
            plot_Isigma.push_back(I_sigma[i] / count[i]);

        if (count[i] == 0 || I[i] < 0.0 || I_sigma[i] < 0.2f * count[i])
            wilson_plot.push_back(NAN);
        else
            wilson_plot.push_back(log(I[i] / count[i]));
    }

    auto res = regression(wilson_plot_legend, wilson_plot);
    if (res.r_square > 0.3)
        return -2.0f * res.slope;

    return {};
}

std::vector<float> BraggIntegrationStats::GetWilsonPlot() {
    return wilson_plot;
}

std::vector<float> BraggIntegrationStats::GetWilsonPlotLegend() {
    return wilson_plot_legend;
}

std::vector<float> BraggIntegrationStats::GetISigmaPlot() {
    return plot_Isigma;
}