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

#include "IceRingVerdict.h"

#include <cmath>

#include "../../common/Definitions.h"
#include "../../common/JFJochMath.h"
#include "../bragg_integration/CalcISigma.h"

namespace {
    struct Moments {
        int n = 0;
        double sum = 0.0, sum2 = 0.0, sum3 = 0.0, sum4 = 0.0;
        void Add(double t) { ++n; sum += t; sum2 += t * t; sum3 += t * t * t; sum4 += t * t * t * t; }
        [[nodiscard]] double Mean() const { return sum / n; }
        [[nodiscard]] double Var() const { return (sum2 - sum * sum / n) / (n - 1); }
        // The variance of ln(Var()): (kurtosis - 1) / n. Wilson intensities are far from Gaussian (an
        // acentric E^2 has kurtosis 9, not 3), so the Gaussian 2 / n would call ordinary scatter excess.
        [[nodiscard]] double VarLogVar() const {
            const double m = Mean();
            const double m2 = sum2 / n - m * m;
            const double m4 = sum4 / n - 4.0 * m * sum3 / n + 6.0 * m * m * sum2 / n - 3.0 * m * m * m * m;
            return (m4 / (m2 * m2) - 1.0) / n;
        }
    };

    double OneOverD2(float d) { return 1.0 / (static_cast<double>(d) * d); }
}

const char *IceRingActionName(IceRingAction a) {
    switch (a) {
        case IceRingAction::DownWeight: return "down-weight";
        case IceRingAction::Reject: return "reject";
        default: return "merge";
    }
}

IceRingDecision DecideIceRings(const std::vector<MergedReflection> &merged) {
    IceRingDecision out;
    out.weight_factor.assign(merged.size(), 1.0);

    std::vector<MergedReflection> off_ring;
    for (const auto &r : merged)
        if (!r.on_ice_ring)
            off_ring.push_back(r);
    const GlobalWilsonB global = CalcGlobalWilsonB(off_ring);

    // The ring each flagged reflection belongs to: the nearest in q.
    std::vector<int> ring_of(merged.size(), -1);
    std::vector<double> dq_of(merged.size(), 0.0);
    for (size_t i = 0; i < merged.size(); ++i) {
        if (!merged[i].on_ice_ring || !(merged[i].d > 0.0f)) continue;
        const double q = 2.0 * PI / merged[i].d;
        for (int k = 0; k < static_cast<int>(ICE_RING_RES_A.size()); ++k) {
            const double dq = q - 2.0 * PI / ICE_RING_RES_A[k];
            if (ring_of[i] < 0 || std::fabs(dq) < std::fabs(dq_of[i])) {
                ring_of[i] = k;
                dq_of[i] = dq;
            }
        }
    }

    for (int k = 0; k < static_cast<int>(ICE_RING_RES_A.size()); ++k) {
        const double s2_ring = OneOverD2(ICE_RING_RES_A[k]);
        bool any = false;
        for (int r : ring_of) any |= r == k;
        if (!any) continue;

        std::vector<MergedReflection> window;
        for (const auto &r : off_ring)
            if (r.d > 0.0f && std::fabs(OneOverD2(r.d) - s2_ring) < ICE_VERDICT_WINDOW_S2)
                window.push_back(r);
        const GlobalWilsonB local = CalcGlobalWilsonB(window);
        const GlobalWilsonB &wilson = local.n_shells >= 3 ? local : global;

        auto t_of = [&](const MergedReflection &r) {
            const double I_exp = std::exp(wilson.intercept + wilson.slope * OneOverD2(r.d));
            return (r.I - I_exp) / std::sqrt(static_cast<double>(r.sigma) * r.sigma + I_exp * I_exp);
        };

        Moments nb;
        if (wilson.n_shells >= 3)
            for (const auto &r : window)
                nb.Add(t_of(r));

        Moments zone[3];
        for (size_t i = 0; i < merged.size(); ++i)
            if (ring_of[i] == k && nb.n > 0)
                zone[dq_of[i] < -ICE_VERDICT_CORE_Q ? 0 : (dq_of[i] > ICE_VERDICT_CORE_Q ? 2 : 1)].Add(t_of(merged[i]));

        IceRingZoneVerdict v[3];
        for (int z = 0; z < 3; ++z) {
            v[z].ring_d = ICE_RING_RES_A[k];
            v[z].zone = z - 1;
            v[z].n = zone[z].n;
            v[z].n_neighbours = nb.n;
            if (zone[z].n < ICE_VERDICT_MIN_ZONE || nb.n < ICE_VERDICT_MIN_NEIGHBOURS)
                continue;
            v[z].bias_z = (zone[z].Mean() - nb.Mean())
                          / std::sqrt(zone[z].Var() / zone[z].n + nb.Var() / nb.n);
            v[z].var_ratio = zone[z].Var() / nb.Var();
            v[z].var_ratio_z = std::log(v[z].var_ratio) / std::sqrt(zone[z].VarLogVar() + nb.VarLogVar());
            if (std::fabs(v[z].bias_z) > ICE_VERDICT_SIGMA)
                v[z].action = IceRingAction::Reject;
            else if (v[z].var_ratio_z > ICE_VERDICT_SIGMA)
                v[z].action = IceRingAction::DownWeight;
        }

        for (size_t i = 0; i < merged.size(); ++i) {
            if (ring_of[i] != k) continue;
            const IceRingZoneVerdict &zv = v[dq_of[i] < -ICE_VERDICT_CORE_Q ? 0 : (dq_of[i] > ICE_VERDICT_CORE_Q ? 2 : 1)];
            if (zv.action == IceRingAction::Reject)
                out.weight_factor[i] = 0.0;
            else if (zv.action == IceRingAction::DownWeight) {
                // The band's excess variance of t, in the reflection's own units.
                const auto &r = merged[i];
                const double I_exp = std::exp(wilson.intercept + wilson.slope * OneOverD2(r.d));
                const double var = static_cast<double>(r.sigma) * r.sigma;
                const double excess = (zone[zv.zone + 1].Var() - nb.Var()) * (var + I_exp * I_exp);
                out.weight_factor[i] = var / (var + excess);
            }
        }
        for (const auto &zv : v)
            if (zv.n > 0)
                out.zones.push_back(zv);
    }
    return out;
}