// 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 <map>
#include <tuple>
#include <vector>

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

namespace {
    MergedReflection Refl(int h, int k, int l, float d) {
        MergedReflection r;
        r.h = h; r.k = k; r.l = l; r.d = d;
        r.I = 100.0f; r.sigma = 10.0f;
        return r;
    }
    // A spread of reflections with a monotone, mate-consistent d (mates share |hkl|).
    std::vector<MergedReflection> Grid(int hmin, int hmax) {
        std::vector<MergedReflection> v;
        for (int h = hmin; h <= hmax; ++h)
            for (int k = 0; k <= 15; ++k)
                for (int l = 0; l <= 15; ++l) {
                    if (h == 0 && k == 0 && l == 0) continue;
                    v.push_back(Refl(h, k, l, 60.0f / (1 + h * h + k * k + l * l)));
                }
        return v;
    }
    double FreeFraction(const std::vector<MergedReflection>& v) {
        int n = 0; for (const auto& r : v) n += r.rfree_flag;
        return static_cast<double>(n) / v.size();
    }
}

TEST_CASE("R-free flags are deterministic and hit the requested fraction", "[rfree]") {
    auto a = Grid(-15, 15);
    auto b = a;
    AssignRfreeFlags(a, 1, 0.05);
    AssignRfreeFlags(b, 1, 0.05);
    REQUIRE(a.size() == b.size());
    for (size_t i = 0; i < a.size(); ++i)
        CHECK(a[i].rfree_flag == b[i].rfree_flag);   // pure function of the reflection

    const double frac = FreeFraction(a);
    CHECK(frac > 0.03);
    CHECK(frac < 0.08);
}

TEST_CASE("R-free flags never split a Friedel/Bijvoet pair", "[rfree]") {
    // Anomalous representation: I(+) and I(-) are separate rows with the same |hkl|.
    std::vector<MergedReflection> v;
    for (int h = 1; h <= 12; ++h)
        for (int k = 0; k <= 12; ++k)
            for (int l = 0; l <= 12; ++l) {
                const float d = 60.0f / (1 + h * h + k * k + l * l);
                v.push_back(Refl(h, k, l, d));
                v.push_back(Refl(-h, -k, -l, d));
            }
    AssignRfreeFlags(v, 1, 0.10);   // P1 -> only Friedel relates the mates

    std::map<std::tuple<int, int, int>, bool> flag;
    for (const auto& r : v) flag[{r.h, r.k, r.l}] = r.rfree_flag;
    int pairs = 0, split = 0;
    for (const auto& r : v) {
        auto it = flag.find({-r.h, -r.k, -r.l});
        if (it != flag.end()) { ++pairs; if (it->second != r.rfree_flag) ++split; }
    }
    CHECK(pairs > 0);
    CHECK(split == 0);
}

TEST_CASE("R-free flags are shared across symmetry equivalents", "[rfree]") {
    // In P4 (Laue 4/m) (h,k,l) and (-k,h,l) are equivalent and must share a flag.
    std::vector<MergedReflection> v;
    for (int h = -10; h <= 10; ++h)
        for (int k = -10; k <= 10; ++k)
            for (int l = 0; l <= 10; ++l) {
                if (h == 0 && k == 0 && l == 0) continue;
                v.push_back(Refl(h, k, l, 60.0f / (1 + h * h + k * k + l * l)));
            }
    AssignRfreeFlags(v, 75, 0.10);   // P4

    std::map<std::tuple<int, int, int>, bool> flag;
    for (const auto& r : v) flag[{r.h, r.k, r.l}] = r.rfree_flag;
    int checked = 0;
    for (const auto& r : v) {
        auto it = flag.find({-r.k, r.h, r.l});   // the 4-fold image
        if (it != flag.end()) { CHECK(it->second == r.rfree_flag); ++checked; }
    }
    CHECK(checked > 0);
}

TEST_CASE("R-free flags are stratified across resolution shells", "[rfree]") {
    // Reflections in three well-separated resolution bands: each band must get some free flags,
    // i.e. the free set is not clumped into one shell.
    std::vector<MergedReflection> v;
    for (int i = 0; i < 1000; ++i) {
        v.push_back(Refl(1 + i, 2, 3, 8.0f));    // low res
        v.push_back(Refl(2, 1 + i, 3, 4.0f));    // mid res
        v.push_back(Refl(2, 3, 1 + i, 2.0f));    // high res
    }
    AssignRfreeFlags(v, 1, 0.10);
    int lo = 0, mid = 0, hi = 0;
    for (const auto& r : v) {
        if (!r.rfree_flag) continue;
        if (r.d > 6.0f) ++lo; else if (r.d > 3.0f) ++mid; else ++hi;
    }
    CHECK(lo > 0);
    CHECK(mid > 0);
    CHECK(hi > 0);
}

TEST_CASE("R-free fraction of zero flags nothing", "[rfree]") {
    auto v = Grid(1, 6);
    AssignRfreeFlags(v, 1, 0.0);
    for (const auto& r : v) CHECK(!r.rfree_flag);
}