// 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_map>
#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);
}

TEST_CASE("Reindex into the ASU: a twin law permutes the reflections without losing any", "[reindex]") {
    const int sg = 75;   // P4
    const auto reference = DistinctReference(sg);
    const auto laws = ReindexAmbiguityOperators(Tetragonal(), sg);
    REQUIRE_FALSE(laws.empty());

    // Mis-index by a twin law, then reindex back into the ASU: the labels must land where an export
    // needs them, and the reflection the label carries must be the one the reference has there.
    const auto misindexed = ReindexReflections(reference, laws.front());
    const auto fixed = ReindexMergedIntoAsu(misindexed, laws.front(), sg, /*merge_friedel=*/true);
    REQUIRE(fixed.size() == reference.size());

    const HKLKeyGenerator key(true, *gemmi::find_spacegroup_by_number(sg));
    std::unordered_map<uint64_t, float> ref_by_key;
    for (const auto &r : reference)
        ref_by_key[key(r).pack()] = r.I;

    std::unordered_set<uint64_t> seen;
    for (const auto &r : fixed) {
        const HKLKey k = key(r);
        CHECK(k.h == r.h);                       // already at its own ASU index
        CHECK(k.k == r.k);
        CHECK(k.l == r.l);
        CHECK(seen.insert(k.pack()).second);     // no two reflections land on one label
        const auto it = ref_by_key.find(k.pack());
        REQUIRE(it != ref_by_key.end());
        CHECK(it->second == r.I);
    }
}

TEST_CASE("Reindex into the ASU: the change of hand swaps the Bijvoet halves", "[reindex]") {
    // The change of hand is the inversion, so it leaves the Laue-ASU label alone and moves the
    // anomalous signal instead - which is the whole of what adopting a model's hand does to intensities.
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(96);   // P4(3)2(1)2
    REQUIRE(sg != nullptr);
    const HKLKeyGenerator key(true, *sg);
    const HKLKey asu = key(3, 1, 2);

    MergedReflection r;
    r.h = asu.h; r.k = asu.k; r.l = asu.l;
    r.I = 100.0f; r.sigma = 1.0f; r.d = 10.0f;
    r.I_plus = 110.0f; r.sigma_plus = 2.0f; r.I_minus = 90.0f; r.sigma_minus = 3.0f;
    r.F_plus = 10.5f; r.F_minus = 9.5f;

    const auto merged = ReindexMergedIntoAsu({r}, sg->change_of_hand_op(), 96, /*merge_friedel=*/true);
    REQUIRE(merged.size() == 1);
    CHECK(merged[0].h == asu.h);
    CHECK(merged[0].k == asu.k);
    CHECK(merged[0].l == asu.l);
    CHECK(merged[0].I == 100.0f);
    CHECK(merged[0].I_plus == 90.0f);
    CHECK(merged[0].I_minus == 110.0f);
    CHECK(merged[0].sigma_plus == 3.0f);
    CHECK(merged[0].sigma_minus == 2.0f);
    CHECK(merged[0].F_plus == 9.5f);
    CHECK(merged[0].F_minus == 10.5f);

    // With the mates kept apart the merge stores the minus hand at -hkl, so the row moves there.
    const auto anom = ReindexMergedIntoAsu({r}, sg->change_of_hand_op(), 96, /*merge_friedel=*/false);
    REQUIRE(anom.size() == 1);
    CHECK(anom[0].h == -asu.h);
    CHECK(anom[0].k == -asu.k);
    CHECK(anom[0].l == -asu.l);
    CHECK(anom[0].I == 100.0f);
}

TEST_CASE("Reindex into the ASU: both mates of an anomalous pair follow the same hand", "[reindex]") {
    // With the mates kept apart, the merge stores the plus mate at +hkl_asu and the minus mate at
    // -hkl_asu, while I_plus/I_minus are attached in the plus convention on BOTH rows alike. An
    // alternative-indexing operator is rotation-type, so op(-x) == -op(x): whatever it does to one
    // mate's hand it does to the other's, and the two rows must still agree afterwards about which
    // intensity is which. Deciding the swap from the hand the new index lands on instead of from the
    // CHANGE of hand moves exactly one of the two, and the pair ends up self-contradictory.
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(75);   // P4, Laue 4/m
    REQUIRE(sg != nullptr);
    const HKLKeyGenerator key_gen(/*merge_friedel=*/false, *sg);
    const HKLKey asu = key_gen(3, 1, 2);
    REQUIRE(asu.plus);

    // The alternative indexing of a P4 lattice: 4/mmm holds it, 4/m does not.
    const gemmi::Op op = gemmi::parse_triplet("k,h,-l");

    MergedReflection plus;
    plus.h = asu.h; plus.k = asu.k; plus.l = asu.l;
    plus.I = 100.0f; plus.sigma = 1.0f; plus.d = 10.0f;
    plus.I_plus = 110.0f; plus.sigma_plus = 2.0f; plus.I_minus = 90.0f; plus.sigma_minus = 3.0f;
    plus.F_plus = 10.5f; plus.F_minus = 9.5f;

    MergedReflection minus = plus;                   // same anomalous split, stored at -hkl_asu
    minus.h = -asu.h; minus.k = -asu.k; minus.l = -asu.l;

    const auto out = ReindexMergedIntoAsu({plus, minus}, op, 75, /*merge_friedel=*/false);
    REQUIRE(out.size() == 2);

    // The two rows are still one pair: same Laue label up to the Friedel sign, opposite hands.
    CHECK(out[0].h == -out[1].h);
    CHECK(out[0].k == -out[1].k);
    CHECK(out[0].l == -out[1].l);

    // ... and they still tell the same story about which hand carries which intensity.
    CHECK(out[0].I_plus == out[1].I_plus);
    CHECK(out[0].I_minus == out[1].I_minus);
    CHECK(out[0].sigma_plus == out[1].sigma_plus);
    CHECK(out[0].sigma_minus == out[1].sigma_minus);
    CHECK(out[0].F_plus == out[1].F_plus);
    CHECK(out[0].F_minus == out[1].F_minus);

    // And the swap did have to happen: "k,h,-l" takes (3,1,2) to (1,3,-2), whose Laue class has its
    // ASU representative at l > 0, so the row that was the plus mate of its pair is the minus mate of
    // the new one. The intensity now standing at the plus index is the one that was at -hkl before.
    CHECK(out[0].I_plus == 90.0f);
    CHECK(out[0].I_minus == 110.0f);
}