// 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 <random>

#include "../common/Definitions.h"
#include "../common/JFJochMath.h"
#include "../image_analysis/scale_merge/IceRingVerdict.h"

namespace {
    // A Wilson crystal (B = 20 A^2) to 2 A, acentric intensities with 10% plus constant noise; the
    // reflections within the flag width of an ice ring are flagged, and `ring` modifies the ones on
    // the 2.667 A ring: dq is their offset from it in q.
    template <class F>
    std::vector<MergedReflection> Crystal(F ring) {
        std::mt19937 rng(1234);
        std::exponential_distribution<double> wilson(1.0);
        std::normal_distribution<double> noise(0.0, 1.0);
        std::uniform_real_distribution<double> s2(1.0 / (6.0 * 6.0), 1.0 / (2.0 * 2.0));
        std::vector<MergedReflection> out;
        for (int i = 0; i < 20000; ++i) {
            MergedReflection r;
            r.d = static_cast<float>(1.0 / std::sqrt(s2(rng)));
            const double I_mean = 1000.0 * std::exp(-0.5 * 20.0 / (r.d * r.d));
            const double I_true = I_mean * wilson(rng);
            r.sigma = static_cast<float>(std::sqrt(0.01 * I_true * I_true + 25.0));
            r.I = static_cast<float>(I_true + r.sigma * noise(rng));
            r.on_ice_ring = IsOnIceRing(r.d, 0.03f);
            const double dq = 2.0 * PI / r.d - 2.0 * PI / 2.667;
            if (std::fabs(dq) < 0.03)
                ring(r, dq, I_mean, rng);
            out.push_back(r);
        }
        return out;
    }

    IceRingAction ActionOf(const IceRingDecision &d, float ring_d, int zone) {
        for (const auto &z : d.zones)
            if (z.ring_d == ring_d && z.zone == zone)
                return z.action;
        return IceRingAction::Merge;
    }
}

TEST_CASE("IceRingVerdict_untouched_ring_is_merged") {
    const auto merged = Crystal([](MergedReflection &, double, double, std::mt19937 &) {});
    const auto d = DecideIceRings(merged);
    REQUIRE(!d.zones.empty());
    for (const auto &z : d.zones)
        CHECK(z.action == IceRingAction::Merge);
    for (double f : d.weight_factor)
        CHECK(f == 1.0);
}

TEST_CASE("IceRingVerdict_biased_core_is_rejected") {
    // The ring adds its own intensity on the core, and takes it out of the flanks' background.
    const auto merged = Crystal([](MergedReflection &r, double dq, double I_mean, std::mt19937 &) {
        r.I += static_cast<float>(std::fabs(dq) < ICE_VERDICT_CORE_Q ? 2.0 * I_mean : -0.5 * I_mean);
    });
    const auto d = DecideIceRings(merged);
    CHECK(ActionOf(d, 2.667f, 0) == IceRingAction::Reject);
    CHECK(ActionOf(d, 2.667f, -1) == IceRingAction::Reject);
    CHECK(ActionOf(d, 2.667f, 1) == IceRingAction::Reject);
    CHECK(ActionOf(d, 2.249f, 0) == IceRingAction::Merge);
    size_t rejected = 0;
    for (double f : d.weight_factor)
        rejected += f == 0.0;
    CHECK(rejected > 0);
}

TEST_CASE("IceRingVerdict_noisy_unbiased_ring_is_down_weighted") {
    // Zero-mean extra scatter of the size of the mean intensity: no bias, only noise.
    const auto merged = Crystal([](MergedReflection &r, double, double I_mean, std::mt19937 &rng) {
        std::normal_distribution<double> extra(0.0, 2.0 * I_mean);
        r.I += static_cast<float>(extra(rng));
    });
    const auto d = DecideIceRings(merged);
    CHECK(ActionOf(d, 2.667f, 0) == IceRingAction::DownWeight);
    for (size_t i = 0; i < merged.size(); ++i)
        if (std::fabs(2.0 * PI / merged[i].d - 2.0 * PI / 2.667) < ICE_VERDICT_CORE_Q) {
            CHECK(d.weight_factor[i] < 1.0);
            CHECK(d.weight_factor[i] > 0.0);
        }
}