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

#include "WilsonOutliers.h"

#include "../../common/ParallelFor.h"

#include <algorithm>
#include <cmath>
#include <numeric>

WilsonOutlierResult WilsonOutliers(const std::vector<WilsonObservation> &obs, double alpha, size_t nthreads) {
    WilsonOutlierResult out;
    out.rejected.assign(obs.size(), 0);
    out.e2.assign(obs.size(), NAN);

    // The observations that can be judged, in index order: fixed blocks count theirs, a prefix places
    // them, and every block writes its own stretch.
    const auto judged = [&](size_t i) {
        const auto &o = obs[i];
        return o.unit >= 0 && o.d > 0.0f && o.d < WILSON_OUTLIER_D_MAX && std::isfinite(o.I)
               && o.sigma > 0.0f && std::isfinite(o.sigma);
    };
    const int n_obs = static_cast<int>(obs.size());
    const int nb = ReductionBlocks(n_obs, 16384);
    const auto block_begin = [&](int b) { return static_cast<size_t>(static_cast<int64_t>(n_obs) * b / nb); };
    std::vector<size_t> block_at(nb + 1, 0);
    ParallelFor(nb, nthreads, [&](int b) {
        size_t c = 0;
        for (size_t i = block_begin(b); i < block_begin(b + 1); ++i) c += judged(i) ? 1 : 0;
        block_at[b + 1] = c;
    });
    for (int b = 0; b < nb; ++b) block_at[b + 1] += block_at[b];
    std::vector<size_t> idx(nb > 0 ? block_at[nb] : 0);
    ParallelFor(nb, nthreads, [&](int b) {
        size_t k = block_at[b];
        for (size_t i = block_begin(b); i < block_begin(b + 1); ++i)
            if (judged(i)) idx[k++] = i;
    });
    if (idx.empty())
        return out;

    // alpha / N per observation, half of it to each tail. Wilson's acentric law is P(E^2 > t) = exp(-t)
    // and the centric one P(E^2 > t) = erfc(sqrt(t/2)) <= exp(-t/2), so t and 2t; the Gaussian tail is
    // P(> z sigma) <= exp(-z^2/2) / 2. Both bounds are the conservative side of the exact quantile.
    // Following Wilson (1949) Acta Cryst. 2, 318-321; the significance guard as in Aimless, Evans (2006)
    // Acta Cryst. D62, 72-82.
    const double p = alpha / (2.0 * static_cast<double>(idx.size()));
    const double t = std::log(1.0 / p);
    const double z = std::sqrt(2.0 * std::log(1.0 / (2.0 * p)));

    // <I/epsilon> in shells of equal observation count, low to high resolution. Two thousand observations
    // put the shell mean at 2% (an exponential's sd equals its mean); narrow shells keep the fall-off
    // across one from reading as spread.
    constexpr size_t OBS_PER_SHELL = 2000;
    // Sorted with the key beside the index, so a comparison reads the two elements it compares rather
    // than two observations somewhere else in memory; the same order.
    {
        struct ByD { float d; size_t i; };
        std::vector<ByD> by_d(idx.size());
        ParallelChunks(static_cast<int>(idx.size()), nthreads, [&](int lo, int hi) {
            for (int j = lo; j < hi; ++j) by_d[j] = {obs[idx[j]].d, idx[j]};
        });
        ParallelSort(by_d.begin(), by_d.end(), nthreads, [](const ByD &a, const ByD &b) {
            return a.d != b.d ? a.d > b.d : a.i < b.i;
        });
        ParallelChunks(static_cast<int>(idx.size()), nthreads, [&](int lo, int hi) {
            for (int j = lo; j < hi; ++j) idx[j] = by_d[j].i;
        });
    }
    const size_t n_shells = std::max<size_t>(1, idx.size() / OBS_PER_SHELL);
    // The observations in that order, so the walks below read them one after another instead of each
    // through idx, and where each shell starts: shell s is the j with j * n_shells / N == s, which begins
    // at the first j with j * n_shells >= s * N. Each shell is then one task, and every sum below is still
    // taken over its shell in the order of j, so the split changes no rounding.
    std::vector<WilsonObservation> so(idx.size());
    ParallelChunks(static_cast<int>(idx.size()), nthreads, [&](int lo, int hi) {
        for (int j = lo; j < hi; ++j) so[j] = obs[idx[j]];
    });
    std::vector<size_t> shell_start(n_shells + 1);
    for (size_t sh = 0; sh <= n_shells; ++sh)
        shell_start[sh] = (sh * idx.size() + n_shells - 1) / n_shells;
    const int n_shell_tasks = static_cast<int>(n_shells);
    const auto lower_e2 = [&](const WilsonObservation &o, double mu) {
        return (o.I - z * o.sigma) / (o.epsilon * mu);
    };
    const auto wilson_bound = [&](const WilsonObservation &o) { return o.centric ? 2.0 * t : t; };
    // A shell is judged only where its mean is itself established at the same significance: with no
    // measured <I>, nothing in the shell can be called improbable against it.
    std::vector<double> mu(n_shells, 0.0);
    std::vector<uint8_t> shell_valid(n_shells, 0);
    ParallelFor(n_shell_tasks, nthreads, [&](int sh) {
        double sum = 0.0;
        size_t cnt = 0;
        for (size_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j) {
            sum += so[j].I / so[j].epsilon;
            ++cnt;
        }
        // Once more without the observations the plain Wilson bound already calls improbable: a single
        // artefact thousands of times <I> would otherwise raise the mean of its own shell several-fold
        // and hide the smaller artefacts sharing it.
        double sum_kept = 0.0, sum2_kept = 0.0;
        size_t cnt_kept = 0;
        for (size_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j) {
            const auto &o = so[j];
            if (sum > 0.0 && lower_e2(o, sum / cnt) > wilson_bound(o))
                continue;
            const double x = o.I / o.epsilon;
            sum_kept += x; sum2_kept += x * x; ++cnt_kept;
        }
        if (cnt_kept < 2) return;
        const double n = static_cast<double>(cnt_kept);
        mu[sh] = sum_kept / n;
        const double var = std::max(0.0, (sum2_kept - n * mu[sh] * mu[sh]) / (n - 1.0));
        shell_valid[sh] = mu[sh] > 0.0 && mu[sh] > z * std::sqrt(var / n);
    });

    // The tail's scale, peaks over threshold: above any threshold an exponential's excess is the same
    // exponential, so the median excess of the acentric observations over a common threshold u is
    // k ln 2. u is where Wilson leaves a hundred acentric observations above it. Only shells whose
    // typical observation is significant at u take part: where sigma is comparable to u <I>, what
    // exceeds u is the noise tail, and it would read as a heavy intensity tail.
    size_t n_acentric = 0;
    for (size_t sh = 0; sh < n_shells; ++sh)
        if (shell_valid[sh])
            for (size_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j)
                if (!so[j].centric) ++n_acentric;
    const double u = std::max(0.0, std::log(static_cast<double>(n_acentric) / 100.0));
    std::vector<uint8_t> shell_measures_tail(n_shells, 0);
    ParallelFor(n_shell_tasks, nthreads, [&](int sh) {
        if (!shell_valid[sh]) return;
        std::vector<double> v;   // sigma / (epsilon <I/epsilon>)
        v.reserve(shell_start[sh + 1] - shell_start[sh]);
        for (size_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j)
            v.push_back(so[j].sigma / (so[j].epsilon * mu[sh]));
        if (v.empty()) return;
        const size_t mid = v.size() / 2;
        std::nth_element(v.begin(), v.begin() + mid, v.end());
        shell_measures_tail[sh] = z * v[mid] <= u;
    });
    std::vector<double> excess;
    for (size_t sh = 0; sh < n_shells; ++sh) {
        if (!shell_measures_tail[sh]) continue;
        for (size_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j) {
            const auto &o = so[j];
            if (o.centric) continue;
            const double e2 = o.I / (o.epsilon * mu[sh]);
            if (e2 > u) excess.push_back(e2 - u);
        }
    }
    if (!excess.empty()) {
        const size_t mid = excess.size() / 2;
        std::nth_element(excess.begin(), excess.begin() + mid, excess.end());
        out.tail_scale = std::max(1.0, excess[mid] / std::log(2.0));
    }
    out.bound = out.tail_scale * t;

    // Which observations are improbable, and the observations of each reflection (in valid shells).
    // Each observation is written once, at its own index, so the shells can do it side by side; the
    // reflections are then formed by a sort on a total order, whatever order tested was filled in.
    std::vector<uint8_t> improbable(obs.size(), 0);
    ParallelFor(n_shell_tasks, nthreads, [&](int sh) {
        if (!shell_valid[sh]) return;
        for (size_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j) {
            const auto &o = so[j];
            out.e2[idx[j]] = static_cast<float>(o.I / (o.epsilon * mu[sh]));
            improbable[idx[j]] = lower_e2(o, mu[sh]) > out.tail_scale * wilson_bound(o);
        }
    });
    std::vector<size_t> tested;
    for (size_t sh = 0; sh < n_shells; ++sh)
        if (shell_valid[sh])
            tested.insert(tested.end(), idx.begin() + shell_start[sh], idx.begin() + shell_start[sh + 1]);
    out.n_tested = tested.size();
    {
        struct ByUnit { int32_t unit; size_t i; };
        std::vector<ByUnit> by_unit(tested.size());
        ParallelChunks(static_cast<int>(tested.size()), nthreads, [&](int lo, int hi) {
            for (int j = lo; j < hi; ++j) by_unit[j] = {obs[tested[j]].unit, tested[j]};
        });
        ParallelSort(by_unit.begin(), by_unit.end(), nthreads, [](const ByUnit &a, const ByUnit &b) {
            return a.unit != b.unit ? a.unit < b.unit : a.i < b.i;
        });
        ParallelChunks(static_cast<int>(tested.size()), nthreads, [&](int lo, int hi) {
            for (int j = lo; j < hi; ++j) tested[j] = by_unit[j].i;
        });
    }

    // An improbable observation measured once is rejected. One with company is rejected when most of
    // its reflection's other observations are probable, unclipped witnesses and it disagrees with their
    // mean beyond the errors - for a pair, the other member. Where most are large, the reflection is
    // (Aimless's "keep if most observations are large").
    // Each reflection is decided on its own observations alone, so the reflections go side by side:
    // where each one's stretch of `tested` starts, then the stretches in blocks.
    const int n_tested = static_cast<int>(tested.size());
    std::vector<uint8_t> starts_unit(n_tested);
    ParallelChunks(n_tested, ThreadsForWork(tested.size(), nthreads), [&](int lo, int hi) {
        for (int q = lo; q < hi; ++q)
            starts_unit[q] = q == 0 || obs[tested[q]].unit != obs[tested[q - 1]].unit;
    });
    std::vector<int32_t> unit_start;
    for (int q = 0; q < n_tested; ++q)
        if (starts_unit[q]) unit_start.push_back(q);
    const int n_units = static_cast<int>(unit_start.size());
    unit_start.push_back(n_tested);
    std::vector<size_t> rejected_in_block(ReductionBlocks(n_units, 1024), 0);
    ParallelBlocks(n_units, nthreads, [&](int blk, int u0, int u1) {
        for (int u = u0; u < u1; ++u) {
            const size_t lo = unit_start[u], hi = unit_start[u + 1];
            for (size_t q = lo; q < hi; ++q) {
                const size_t i = tested[q];
                if (!improbable[i]) continue;
                size_t n_probable = 0;
                double sw = 0.0, swI = 0.0;
                for (size_t r = lo; r < hi; ++r) {
                    const auto &o = obs[tested[r]];
                    if (r == q || improbable[tested[r]] || o.clipped) continue;
                    const double w = 1.0 / (static_cast<double>(o.sigma) * o.sigma);
                    sw += w; swI += w * o.I;
                    ++n_probable;
                }
                const size_t n_others = hi - lo - 1;
                bool reject = n_others == 0;
                if (2 * n_probable > n_others) {
                    const auto &o = obs[i];
                    reject = std::fabs(o.I - swI / sw) > z * std::sqrt(static_cast<double>(o.sigma) * o.sigma + 1.0 / sw);
                }
                if (reject) {
                    out.rejected[i] = 1;
                    ++rejected_in_block[blk];
                }
            }
        }
    }, 1024);
    for (const size_t r : rejected_in_block) out.n_rejected += r;
    return out;
}