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

#include <catch2/catch_all.hpp>

#include <cmath>
#include <random>
#include <vector>

#include "../image_analysis/scale_merge/WilsonOutliers.h"

namespace {
    // Acentric Wilson intensities of mean `mean` between 3.9 and 2.0 A, measured `per_unit` times each
    // with Poisson-like noise of variance I + bkg_var.
    std::vector<WilsonObservation> WilsonPopulation(int n_units, int per_unit, double mean, double bkg_var,
                                                    uint32_t seed) {
        std::mt19937 rng(seed);
        std::exponential_distribution<double> wilson(1.0 / mean);
        std::normal_distribution<double> gauss(0.0, 1.0);
        std::vector<WilsonObservation> v;
        for (int u = 0; u < n_units; ++u) {
            const double I_true = wilson(rng);
            const float d = 3.9f - 1.9f * static_cast<float>(u) / n_units;
            for (int m = 0; m < per_unit; ++m) {
                const double sigma = std::sqrt(I_true + bkg_var);
                v.push_back({static_cast<float>(I_true + sigma * gauss(rng)), static_cast<float>(sigma), d,
                             1.0f, false, false, u});
            }
        }
        return v;
    }
}

TEST_CASE("WilsonOutliers: the artefact member of a discordant pair is dropped", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 1);
    // A pair of one ordinary observation and one two hundred times the shell mean.
    const int32_t unit = 10000;
    obs.push_back({800.0f, 30.0f, 2.5f, 1.0f, false, false, unit});
    obs.push_back({200000.0f, 450.0f, 2.5f, 1.0f, false, false, unit});

    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == obs.size());
    CHECK(r.tail_scale == Catch::Approx(1.0).margin(0.2));
    CHECK(r.rejected[obs.size() - 1] == 1);
    CHECK(r.rejected[obs.size() - 2] == 0);
    CHECK(r.n_rejected == 1);
}

TEST_CASE("WilsonOutliers: a clipped mate does not testify against a strong observation", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 10);
    // The strong member is real; the low one lost its saturated core to the mask.
    const int32_t unit = 10000;
    obs.push_back({800.0f, 30.0f, 2.5f, 1.0f, false, true, unit});
    obs.push_back({200000.0f, 450.0f, 2.5f, 1.0f, false, false, unit});
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: two large observations of one reflection confirm each other", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 2);
    const int32_t unit = 10000;
    obs.push_back({60000.0f, 300.0f, 2.5f, 1.0f, false, false, unit});
    obs.push_back({61000.0f, 300.0f, 2.5f, 1.0f, false, false, unit});

    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: the symmetry enhancement factor keeps an axial reflection", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 3);
    // Measured once, 12x <I/epsilon> at epsilon 4: 48x the shell mean, but only E^2 = 12.
    obs.push_back({48000.0f, 250.0f, 2.5f, 4.0f, false, false, 10000});
    auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.rejected.back() == 0);
    CHECK(r.e2.back() == Catch::Approx(12.0).epsilon(0.1));
    CHECK(r.n_rejected == 0);

    // The same observation on a general reflection is improbable.
    obs.back().epsilon = 1.0f;
    r = WilsonOutliers(obs, 0.01);
    CHECK(r.rejected.back() == 1);
}

TEST_CASE("WilsonOutliers: a weak shell does not misfire", "[wilson_outliers]") {
    // <I> = 20 under a background of sd 100: noise alone reaches twenty times <I>, while the shell mean
    // is still established.
    const auto obs = WilsonPopulation(20000, 1, 20.0, 10000.0, 4);
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == obs.size());
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: an artefact among several mates is dropped", "[wilson_outliers]") {
    auto obs = WilsonPopulation(5000, 4, 1000.0, 100.0, 5);
    // One precise-looking artefact beside three ordinary mates: the mates out-vote it.
    obs[0].I = 500000.0f;
    obs[0].sigma = 700.0f;
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == obs.size());
    CHECK(r.rejected[0] == 1);
    CHECK(r.n_rejected == 1);
}

TEST_CASE("WilsonOutliers: a reflection whose observations are mostly large is kept", "[wilson_outliers]") {
    auto obs = WilsonPopulation(5000, 3, 1000.0, 100.0, 8);
    // Two of three observations large, the third low (a partial that caught little, say).
    obs[0].I = 60000.0f; obs[0].sigma = 300.0f;
    obs[1].I = 62000.0f; obs[1].sigma = 300.0f;
    obs[2].I = 500.0f; obs[2].sigma = 30.0f;
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: a shell without a measured mean is not judged", "[wilson_outliers]") {
    // Pure noise: <I> = 0 within its error, so nothing can be improbable against it.
    auto obs = WilsonPopulation(10000, 1, 1e-6, 10000.0, 9);
    obs[0].I = 5000.0f; obs[0].sigma = 100.0f;
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == 0);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: a heavier-tailed population widens the bound", "[wilson_outliers]") {
    // Half the reflections at twice the mean, half at a tenth - the intensity classes of a strong
    // pseudo-translation. Wilson's single exponential would call the top of the strong class improbable.
    auto strong = WilsonPopulation(10000, 1, 2000.0, 100.0, 6);
    const auto weak = WilsonPopulation(10000, 1, 100.0, 100.0, 7);
    for (auto o : weak) {
        o.unit += 10000;
        strong.push_back(o);
    }
    const auto r = WilsonOutliers(strong, 0.01);
    CHECK(r.tail_scale > 1.5);
    CHECK(r.n_rejected == 0);
}