// 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;
    // Floor off (min_free = 0) so this isolates the pure-hash fraction on this modest grid.
    AssignRfreeFlags(a, 1, 0.05, /*min_free=*/0);
    AssignRfreeFlags(b, 1, 0.05, /*min_free=*/0);
    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 spread across resolution", "[rfree]") {
    // The uniform per-hkl hash is uncorrelated with resolution, so each of three well-separated
    // resolution bands still receives a share of the free set (it 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 flags are identical across datasets of one crystal form", "[rfree]") {
    // The key campaign property: the free set depends only on the reflection index, not on the
    // dataset's resolution extent or which reflections it contains. Two datasets with different
    // resolution ranges must flag every shared reflection the same way (a per-shell stratification
    // tied to each dataset's own d_min/d_max would break this).
    auto wide = Grid(-15, 15);
    std::vector<MergedReflection> narrow;                 // a lower-resolution subset
    for (const auto& r : wide)
        if (r.d > 6.0f)
            narrow.push_back(r);
    // Floor off so the two different-sized sets share one effective fraction (the floor is the only
    // thing that ties the fraction to the dataset; with it off this is the pure per-hkl guarantee).
    AssignRfreeFlags(wide, 96, 0.05, /*min_free=*/0);
    AssignRfreeFlags(narrow, 96, 0.05, /*min_free=*/0);

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

TEST_CASE("ApplyReferenceFreeFlags imports the reference test set", "[rfree]") {
    // A reference with its own free set; a dataset that starts from the per-hkl hash must, after the
    // import, carry exactly the reference's flags on every reflection they share.
    auto reference = Grid(-12, 12);
    AssignRfreeFlags(reference, 96, 0.07);
    std::map<std::tuple<int, int, int>, bool> ref_flag;
    for (const auto& r : reference) ref_flag[{r.h, r.k, r.l}] = r.rfree_flag;

    auto data = Grid(-12, 12);
    AssignRfreeFlags(data, 96, 0.30);                     // deliberately a different fraction/hash split
    const size_t matched = ApplyReferenceFreeFlags(data, 96, reference);
    CHECK(matched == data.size());
    for (const auto& r : data)
        CHECK(r.rfree_flag == ref_flag[{r.h, r.k, r.l}]);
}

TEST_CASE("R-free flags floor the test set size on small data", "[rfree]") {
    // A small dataset (~738 distinct P1 reflections, all l=+1 so no Friedel mates present): at 5%
    // only ~37 would be free, too few for a stable R-free. The floor lifts the fraction; a large
    // floor request is capped so it never dominates the working set.
    std::vector<MergedReflection> v;
    for (int h = -20; h <= 20; ++h)
        for (int k = 1; k <= 18; ++k)
            v.push_back(Refl(h, k, 1, 30.0f / (1 + h * h + k * k)));
    const double n = static_cast<double>(v.size());

    auto plain = v, floored = v, capped = v;
    AssignRfreeFlags(plain, 1, 0.05, /*min_free=*/0);       // pure 5% ~= 37 free
    AssignRfreeFlags(floored, 1, 0.05, /*min_free=*/60);    // 60/738 = 8.1% (above 5%, under the 10% cap)
    AssignRfreeFlags(capped, 1, 0.05, /*min_free=*/100000); // floor wants ~all; capped near 10%

    CHECK(FreeFraction(floored) > FreeFraction(plain));     // the floor lifted the test set
    CHECK(FreeFraction(floored) > 0.06);                    // ~8%, clearly above the bare 5%
    CHECK(FreeFraction(capped) <= 0.15);                    // capped near 10%, not driven to ~100%
}

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);
}