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

#include "../image_analysis/scale_merge/HKLKey.h"
#include "../image_analysis/scale_merge/ReindexAmbiguity.h"

namespace {
    UnitCell Tetragonal() { return UnitCell{78.0f, 78.0f, 37.0f, 90.0f, 90.0f, 90.0f}; }

    // A reference set: one distinct intensity per Laue-ASU reflection, so that reflections related by
    // a twin law (which are NOT Laue-equivalent in P4) carry different intensities.
    std::vector<MergedReflection> DistinctReference(int space_group_number) {
        const HKLKeyGenerator key(true, *gemmi::find_spacegroup_by_number(space_group_number));
        std::vector<MergedReflection> v;
        std::unordered_set<uint64_t> seen;
        for (int h = -8; h <= 8; ++h)
            for (int k = -8; k <= 8; ++k)
                for (int l = 0; l <= 8; ++l) {
                    if (h == 0 && k == 0 && l == 0) continue;
                    const uint64_t kk = key(h, k, l).pack();
                    if (!seen.insert(kk).second) continue;   // one row per ASU reflection
                    MergedReflection r;
                    r.h = h; r.k = k; r.l = l;
                    r.I = static_cast<float>(100 + kk % 9973);   // distinct, spread
                    r.sigma = 1.0f;
                    r.d = 60.0f / (1 + h * h + k * k + l * l);
                    v.push_back(r);
                }
        return v;
    }
}

TEST_CASE("Reindex: holohedral crystal has no indexing ambiguity", "[reindex]") {
    // P4(3)2(1)2 (422, holohedral for the tetragonal lattice) -> no twin laws.
    CHECK(ReindexAmbiguityOperators(Tetragonal(), 96).empty());
    // P422 likewise.
    CHECK(ReindexAmbiguityOperators(Tetragonal(), 89).empty());
}

TEST_CASE("Reindex: merohedral crystal exposes the ambiguity operators", "[reindex]") {
    // P4 (point group 4) in a tetragonal lattice (422) -> a non-trivial reindexing coset.
    CHECK_FALSE(ReindexAmbiguityOperators(Tetragonal(), 75).empty());
}

TEST_CASE("Reindex: reference agreement recovers a misindexed dataset", "[reindex]") {
    const int sg = 75;   // P4
    const auto reference = DistinctReference(sg);
    const auto laws = ReindexAmbiguityOperators(Tetragonal(), sg);
    REQUIRE_FALSE(laws.empty());

    // Deliberately mis-index the data by one twin law.
    const auto data = ReindexReflections(reference, laws.front());

    const auto choice = ChooseReindex(
        data, Tetragonal(), sg,
        [&](const std::vector<MergedReflection> &m) { return ReferenceIntensityCC(m, reference, sg); });

    CHECK_FALSE(choice.is_identity);                 // a reindex was needed
    CHECK(choice.score > 0.99);                      // the winner realigns with the reference
    CHECK(choice.identity_score < 0.9);              // leaving it mis-indexed correlates poorly
    CHECK(choice.n_candidates >= 2);
}

TEST_CASE("Reindex: a correctly indexed dataset keeps identity", "[reindex]") {
    const int sg = 75;
    const auto reference = DistinctReference(sg);
    const auto choice = ChooseReindex(
        reference, Tetragonal(), sg,
        [&](const std::vector<MergedReflection> &m) { return ReferenceIntensityCC(m, reference, sg); });
    CHECK(choice.is_identity);
    CHECK(choice.score > 0.99);
}