// 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 <algorithm>
#include <array>
#include <cmath>
#include <random>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "../image_analysis/scale_merge/CrystalSetting.h"
#include "../image_analysis/scale_merge/HKLKey.h"
#include "../image_analysis/scale_merge/ReindexAmbiguity.h"
#include "../image_analysis/scale_merge/ReferenceFreeReindex.h"
#include "../image_analysis/scale_merge/SublatticeRelabel.h"

namespace {
    const gemmi::SpaceGroup &SG(int number) { return *gemmi::find_spacegroup_by_number(number); }

    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(), SG(96)).empty());
    // P422 likewise.
    CHECK(ReindexAmbiguityOperators(Tetragonal(), SG(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(), SG(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(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(sg),
        [&](const std::vector<MergedReflection> &m) { return ReferenceIntensityCC(m, reference, SG(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(sg),
        [&](const std::vector<MergedReflection> &m) { return ReferenceIntensityCC(m, reference, SG(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(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(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(), SG(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(), SG(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, SG(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);
}

TEST_CASE("Reindex: every ambiguity operator is a proper rotation", "[reindex]") {
    // A relabelling by a lattice rotation keeps the hand; one with determinant -1 would swap every
    // Bijvoet pair. The candidates come from the lattice symmetry without the inversion, so none may
    // be improper, in any crystal system that has an ambiguity.
    const std::vector<std::pair<UnitCell, int>> cases{
        {Tetragonal(), 75},                                           // P4
        {UnitCell{60.0f, 60.0f, 90.0f, 90.0f, 90.0f, 120.0f}, 143},   // P3
        {UnitCell{60.0f, 60.0f, 90.0f, 90.0f, 90.0f, 120.0f}, 149},   // P312
        {UnitCell{60.0f, 60.0f, 90.0f, 90.0f, 90.0f, 120.0f}, 168},   // P6
        {UnitCell{77.0f, 77.0f, 77.0f, 90.0f, 90.0f, 90.0f}, 197},    // I23
        {UnitCell{77.0f, 77.0f, 77.0f, 90.0f, 90.0f, 90.0f}, 199},    // I2(1)3
    };
    for (const auto &[cell, sg] : cases) {
        const auto ops = ReindexAmbiguityOperators(cell, SG(sg));
        CHECK_FALSE(ops.empty());
        for (const auto &op : ops)
            CHECK(op.det_rot() == gemmi::Op::DEN * gemmi::Op::DEN * gemmi::Op::DEN);
    }
}

TEST_CASE("Reindex: a near-symmetry law moves the cell with the indices", "[reindex]") {
    // Pseudo-merohedry: b within 2% of c in a primitive orthorhombic cell, and beta a degree off 90 in
    // a monoclinic one. Each admits an indexing law that is not a symmetry of the cell, so relabelling
    // the reflections by it permutes the axes (or flips beta), and the cell they are written with has
    // to follow - every reflection then keeps its own d-spacing.
    const std::vector<std::pair<UnitCell, int>> cases{
        {UnitCell{50.0f, 60.0f, 61.2f, 90.0f, 90.0f, 90.0f}, 16},    // P222, b ~ c
        {UnitCell{70.0f, 79.0f, 83.0f, 90.0f, 91.1f, 90.0f}, 3},     // P2, beta ~ 90
    };
    for (const auto &[cell, sg] : cases) {
        const auto laws = ReindexAmbiguityOperators(cell, SG(sg));
        REQUIRE_FALSE(laws.empty());
        bool moved = false;
        for (const gemmi::Op &law : laws) {
            const UnitCell written = CellInBasis(cell, HklOperator(law));
            moved = moved || std::abs(written.b - cell.b) > 0.5f || std::abs(written.beta - cell.beta) > 0.5f;
            const gemmi::UnitCell before = cell, after = written;
            for (const gemmi::Op::Miller hkl : {gemmi::Op::Miller{1, 2, 3}, gemmi::Op::Miller{-4, 0, 7},
                                                gemmi::Op::Miller{5, -3, 1}})
                CHECK(after.calculate_d(law.apply_to_hkl(hkl)) == Catch::Approx(before.calculate_d(hkl)).epsilon(1e-5));
        }
        CHECK(moved);
    }
}

namespace {
    // A P2(1) cell with a near-symmetry: a+c is within 1% of a, so the lattice is metrically close to
    // a C-centred orthorhombic one (obliquity ~0.7 deg) without being one.
    UnitCell PseudoMonoclinic() { return UnitCell{64.0f, 29.0f, 36.0f, 90.0f, 107.0f, 90.0f}; }

    Coord Recip(const CrystalLattice &latt, int h, int k, int l) {
        return latt.Astar() * static_cast<float>(h) + latt.Bstar() * static_cast<float>(k)
             + latt.Cstar() * static_cast<float>(l);
    }
}

TEST_CASE("Reindex: an ambiguity operator of a pseudo-symmetric monoclinic cell keeps every reflection's resolution",
          "[reindex]") {
    const UnitCell cell = PseudoMonoclinic();
    const auto laws = ReindexAmbiguityOperators(cell, SG(4));
    REQUIRE_FALSE(laws.empty());
    const CrystalLattice latt(cell);
    for (const auto &op : laws)
        for (int h = -6; h <= 6; ++h)
            for (int l = -6; l <= 6; ++l) {
                if (h == 0 && l == 0) continue;
                // As the resolver applies it: an operator that is a near-symmetry of the metric moves an
                // index to one of nearly the same length, never to a finer one.
                const auto m = op.apply_to_hkl({{h, 3, l}});
                const float d0 = 1.0f / Recip(latt, h, 3, l).Length();
                const float d1 = 1.0f / Recip(latt, m[0], m[1], m[2]).Length();
                CHECK(std::fabs(d1 - d0) / d0 < 0.03f);
            }
}

TEST_CASE("Reindex: the reindexed lattice predicts every spot where it was", "[reindex]") {
    const UnitCell cell = PseudoMonoclinic();
    const auto laws = ReindexAmbiguityOperators(cell, SG(4));
    REQUIRE_FALSE(laws.empty());
    const CrystalLattice latt(cell);
    for (const auto &op : laws) {
        const CrystalLattice moved = ReindexLattice(latt, op);
        for (int h = -4; h <= 4; ++h)
            for (int k = -4; k <= 4; ++k)
                for (int l = -4; l <= 4; ++l) {
                    const auto m = op.apply_to_hkl({{h, k, l}});
                    const Coord before = Recip(latt, h, k, l);
                    const Coord after = Recip(moved, m[0], m[1], m[2]);
                    CHECK((after - before).Length() < 1e-5f);
                }
    }
}

TEST_CASE("Sublattice: a frame on the given cell in another setting is described on the given cell's axes",
          "[reindex]") {
    // The given C222(1)-style cell, and the same lattice with b and c swapped (and a flipped to keep the
    // basis right-handed): not close axis by axis, but the same lattice.
    const UnitCell given{40.0f, 180.0f, 143.0f, 90.0f, 90.0f, 90.0f};
    const UnitCell frame{40.2f, 143.4f, 180.5f, 90.0f, 90.0f, 90.0f};
    IndexingSettings settings;
    REQUIRE_FALSE(frame.is_close(given, settings.GetUnitCellDistTolerance(), settings.GetUnitCellAngleTolerance_deg()));
    const auto on_given = CellOnReference(frame, given, settings);
    REQUIRE(on_given.has_value());
    CHECK(on_given->is_close(given, settings.GetUnitCellDistTolerance(), settings.GetUnitCellAngleTolerance_deg()));
    CHECK(on_given->b == Catch::Approx(180.5f).epsilon(0.001));
    CHECK(on_given->c == Catch::Approx(143.4f).epsilon(0.001));

    // An unrelated cell is not put on it.
    CHECK_FALSE(CellOnReference(UnitCell{55.0f, 66.0f, 77.0f, 90.0f, 90.0f, 90.0f}, given, settings).has_value());
}

TEST_CASE("Reindex: a stills set indexed in random hands is made consistent from the data alone", "[reindex]") {
    // P4 on a tetragonal lattice: two indexings, related by a twin law. Every crystal records a random
    // part of the same synthetic structure with its own scale and noise, and half of them, at random,
    // carry the other indexing.
    const UnitCell cell = Tetragonal();
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(75);
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 1);

    const HKLKeyGenerator key(true, x.GetSpaceGroupOrP1());
    std::mt19937 rng(7);
    std::exponential_distribution<double> wilson(1.0);
    std::normal_distribution<double> noise(0.0, 0.3);
    std::unordered_map<uint64_t, double> structure;
    std::vector<std::array<int, 3>> hkls;
    for (int h = -12; h <= 12; ++h)
        for (int k = -12; k <= 12; ++k)
            for (int l = -6; l <= 6; ++l) {
                if (h == 0 && k == 0 && l == 0)
                    continue;
                hkls.push_back({h, k, l});
                structure.try_emplace(key(h, k, l).pack(), wilson(rng));
            }

    const int n_crystals = 300;
    std::vector<IntegrationOutcome> outcomes(n_crystals);
    std::vector<bool> flipped(n_crystals);
    for (int c = 0; c < n_crystals; ++c) {
        flipped[c] = rng() % 2;
        const double scale = 0.5 + (rng() % 1000) / 1000.0;
        outcomes[c].latt = CrystalLattice(cell);
        for (int i = 0; i < 150; ++i) {
            const auto &t = hkls[rng() % hkls.size()];
            const auto m = flipped[c] ? laws[0].apply_to_hkl({{t[0], t[1], t[2]}}) : gemmi::Op::Miller{{t[0], t[1], t[2]}};
            Reflection r{};
            r.h = m[0]; r.k = m[1]; r.l = m[2];
            r.d = 1.0f / std::sqrt((t[0] * t[0] + t[1] * t[1]) / (cell.a * cell.a) + t[2] * t[2] / (cell.c * cell.c));
            r.I = static_cast<float>(scale * structure.at(key(t[0], t[1], t[2]).pack()) * (1.0 + noise(rng)));
            r.sigma = 1.0f;
            r.partiality = 1.0f;
            r.prescaling_corr = 1.0f;
            outcomes[c].reflections.push_back(r);
        }
    }

    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 4);
    CHECK(result.crystals == n_crystals);
    // The overall hand is free: either exactly the flipped crystals are reindexed, or exactly the others.
    int agree = 0;
    for (int c = 0; c < n_crystals; ++c)
        agree += (result.choice[c] != 0) == flipped[c];
    CHECK((agree == n_crystals || agree == 0));
    const size_t n_flipped = std::count(flipped.begin(), flipped.end(), true);
    CHECK(result.reindexed == std::min(n_flipped, n_crystals - n_flipped));
    CHECK(result.reindexed <= n_crystals / 2);   // the overall hand that moves the fewest crystals
    CHECK(result.decided_fraction > 0.9);

    // Done again on the consistent set, nothing moves.
    const auto again = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 4);
    CHECK(again.reindexed == 0);
}

TEST_CASE("Reindex: many weak stills with heavy-tailed intensities are made consistent from the data alone",
          "[reindex]") {
    // As above, but in the regime of a serial run: many crystals with few reflections each, a
    // multiplicative error on every intensity as large as the intensity (partiality), a resolution
    // fall-off, and twin-related reflections that often carry similar intensities. Compared on the
    // intensities themselves, a few very strong reflections dominate and a tenth of the crystals end up
    // in the wrong indexing.
    const UnitCell cell = Tetragonal();
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(75);
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 1);

    const HKLKeyGenerator key(true, x.GetSpaceGroupOrP1());
    std::mt19937 rng(5);
    std::exponential_distribution<double> wilson(1.0);
    std::uniform_real_distribution<double> uniform(0.0, 1.0);
    std::normal_distribution<double> normal(0.0, 1.0);
    std::unordered_map<uint64_t, double> structure;
    std::vector<std::array<int, 3>> hkls;
    for (int h = -14; h <= 14; ++h)
        for (int k = -14; k <= 14; ++k)
            for (int l = -8; l <= 8; ++l)
                if (h != 0 || k != 0 || l != 0) {
                    hkls.push_back({h, k, l});
                    structure.try_emplace(key(h, k, l).pack(), wilson(rng));
                }
    // Two twin pairs in five carry nearly the same intensity.
    std::unordered_set<uint64_t> paired;
    for (const auto &t : hkls) {
        const auto m = laws[0].apply_to_hkl({{t[0], t[1], t[2]}});
        const uint64_t a = key(t[0], t[1], t[2]).pack(), b = key(m[0], m[1], m[2]).pack();
        if (a == b || !paired.insert(a).second || !paired.insert(b).second)
            continue;
        if (uniform(rng) < 0.4)
            structure[b] = structure[a] * (0.7 + 0.6 * uniform(rng));
    }

    const int n_crystals = 2000;
    std::vector<IntegrationOutcome> outcomes(n_crystals);
    std::vector<bool> flipped(n_crystals);
    for (int c = 0; c < n_crystals; ++c) {
        flipped[c] = rng() % 2;
        const double scale = std::exp(0.5 * normal(rng));
        outcomes[c].latt = CrystalLattice(cell);
        for (int i = 0; i < 60; ++i) {
            const auto &t = hkls[rng() % hkls.size()];
            const double s2 = (t[0] * t[0] + t[1] * t[1]) / (cell.a * cell.a) + t[2] * t[2] / (cell.c * cell.c);
            const auto m = flipped[c] ? laws[0].apply_to_hkl({{t[0], t[1], t[2]}}) : gemmi::Op::Miller{{t[0], t[1], t[2]}};
            Reflection r{};
            r.h = m[0]; r.k = m[1]; r.l = m[2];
            r.d = static_cast<float>(1.0 / std::sqrt(s2));
            r.I = static_cast<float>(scale * structure.at(key(t[0], t[1], t[2]).pack()) * std::exp(-15.0 * s2)
                                     * std::exp(normal(rng)) + 0.02 * normal(rng));
            r.sigma = 0.02f;
            r.partiality = 1.0f;
            r.prescaling_corr = 1.0f;
            outcomes[c].reflections.push_back(r);
        }
    }

    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 4);
    int agree = 0;
    for (int c = 0; c < n_crystals; ++c)
        agree += (result.choice[c] != 0) == flipped[c];
    CHECK(std::max(agree, n_crystals - agree) > 0.97 * n_crystals);
}

TEST_CASE("Reindex: a holohedral crystal gives the stills resolution nothing to do", "[reindex]") {
    DiffractionExperiment x;
    x.SetUnitCell(Tetragonal()).SpaceGroupNumber(89);   // P422
    std::vector<IntegrationOutcome> outcomes(3);
    const auto laws = ReindexAmbiguityOperators(Tetragonal(), x.GetSpaceGroupOrP1());
    const auto result = ResolveIndexingAmbiguity(outcomes, x, Tetragonal(), laws, {}, 1);
    CHECK(laws.empty());
    CHECK(result.reindexed == 0);
    CHECK(result.crystals == 0);
}

namespace {
    // A crystal's refined cell: `cell` with each length off by a relative error of `scatter`.
    UnitCell Scattered(const UnitCell &cell, double scatter, std::mt19937 &rng) {
        std::normal_distribution<double> e(0.0, scatter);
        return UnitCell{static_cast<float>(cell.a * (1.0 + e(rng))), static_cast<float>(cell.b * (1.0 + e(rng))),
                        static_cast<float>(cell.c * (1.0 + e(rng))), cell.alpha, cell.beta, cell.gamma};
    }

    // A synthetic P222 structure on hkl up to (12, 12, 6), one Wilson-distributed intensity per reflection
    // of the asymmetric unit.
    struct Structure {
        std::unordered_map<uint64_t, double> intensity;
        std::vector<std::array<int, 3>> hkls;
    };
    Structure MakeStructure(const HKLKeyGenerator &key, std::mt19937 &rng) {
        std::exponential_distribution<double> wilson(1.0);
        Structure s;
        for (int h = -12; h <= 12; ++h)
            for (int k = -12; k <= 12; ++k)
                for (int l = -6; l <= 6; ++l)
                    if (h != 0 || k != 0 || l != 0) {
                        s.hkls.push_back({h, k, l});
                        s.intensity.try_emplace(key(h, k, l).pack(), wilson(rng));
                    }
        return s;
    }

    // A crystal of `structure` truly in the other indexing (by `law`) where `flipped`, recording 150 random
    // reflections with a relative noise of 0.3 - in its own reported indices - with the lattice `latt`.
    IntegrationOutcome Crystal(const Structure &structure, const HKLKeyGenerator &key, const UnitCell &cell,
                               const gemmi::Op &law, bool flipped, const CrystalLattice &latt, std::mt19937 &rng) {
        std::normal_distribution<double> noise(0.0, 0.3);
        IntegrationOutcome o;
        o.latt = latt;
        for (int i = 0; i < 150; ++i) {
            const auto &t = structure.hkls[rng() % structure.hkls.size()];
            const auto m = flipped ? law.apply_to_hkl({{t[0], t[1], t[2]}}) : gemmi::Op::Miller{{t[0], t[1], t[2]}};
            Reflection r{};
            r.h = m[0]; r.k = m[1]; r.l = m[2];
            r.d = 1.0f / std::sqrt(t[0] * t[0] / (cell.a * cell.a) + t[1] * t[1] / (cell.b * cell.b)
                                   + t[2] * t[2] / (cell.c * cell.c));
            r.I = static_cast<float>(structure.intensity.at(key(t[0], t[1], t[2]).pack()) * (1.0 + noise(rng)));
            r.sigma = 0.1f;
            r.partiality = 1.0f;
            r.prescaling_corr = 1.0f;
            o.reflections.push_back(r);
        }
        return o;
    }
}

TEST_CASE("Reindex: a crystal's own cell says which axis is which where the run's axes differ beyond its scatter",
          "[reindex]") {
    // Orthorhombic with a and b 3% apart: a twin law of the metric to within the obliquity tolerance, but each
    // crystal's refined cell, scattering by 0.3%, says which axis is which. The intensities say nothing.
    const UnitCell cell{78.0f, 80.4f, 37.0f, 90.0f, 90.0f, 90.0f};
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(16);   // P222
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 1);

    std::vector<IntegrationOutcome> outcomes(50);
    std::mt19937 rng(3);
    for (auto &o : outcomes) {
        o.latt = CrystalLattice(Scattered(cell, 0.003, rng));
        for (int i = 0; i < 100; ++i) {
            Reflection r{};
            r.h = static_cast<int>(rng() % 15) - 7; r.k = static_cast<int>(rng() % 15) - 7; r.l = static_cast<int>(rng() % 7) + 1;
            r.d = 3.0f; r.I = static_cast<float>(rng() % 1000); r.sigma = 1.0f; r.partiality = 1.0f; r.prescaling_corr = 1.0f;
            o.reflections.push_back(r);
        }
    }
    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 1);
    REQUIRE(result.metric_log_odds.size() == 1);
    CHECK(result.metric_log_odds[0] > 10.0);
    CHECK(result.model.metric.variance[0] == Catch::Approx(0.003 * 0.003).epsilon(0.5));
    CHECK(result.crystals == 50);
    CHECK(result.reindexed == 0);
}

TEST_CASE("Reindex: where the cells tell the axes apart, they and the intensities put back the crystals indexed on swapped axes",
          "[reindex]") {
    // As above, a and b 3% apart, but a few crystals were indexed on swapped axes: their own cells say so,
    // and so do their intensities.
    const UnitCell cell{78.0f, 80.4f, 37.0f, 90.0f, 90.0f, 90.0f};
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(16);   // P222
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 1);

    const HKLKeyGenerator key(true, x.GetSpaceGroupOrP1());
    std::mt19937 rng(11);
    const Structure structure = MakeStructure(key, rng);
    const int n_crystals = 300;
    std::vector<IntegrationOutcome> outcomes;
    std::vector<bool> swapped(n_crystals, false);
    for (int c = 0; c < n_crystals; ++c) {
        swapped[c] = c % 12 == 5;
        const CrystalLattice latt(Scattered(cell, 0.003, rng));
        outcomes.push_back(Crystal(structure, key, cell, laws[0], swapped[c],
                                   swapped[c] ? ReindexLattice(latt, laws[0]) : latt, rng));
    }

    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 4);
    CHECK(result.crystals == n_crystals);
    for (int c = 0; c < n_crystals; ++c)
        CHECK((result.choice[c] != 0) == swapped[c]);
}

TEST_CASE("Reindex: cells that only echo the hand the indexer chose do not lock the crystals in it", "[reindex]") {
    // a and b 0.2% apart, well inside the 1% the crystals' refined cells scatter by. Every crystal is truly
    // in one of the two indexings at random, but its cell is reported in whichever setting fits the run's
    // cell better - as the indexer or a setting convention picks it - so every cell prefers the indexing it
    // came in, whatever the truth. The intensities have to decide, and do.
    const UnitCell cell{78.0f, 78.16f, 37.0f, 90.0f, 90.0f, 90.0f};
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(16);   // P222
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 1);

    const HKLKeyGenerator key(true, x.GetSpaceGroupOrP1());
    std::mt19937 rng(17);
    const Structure structure = MakeStructure(key, rng);
    const int n_crystals = 300;
    std::vector<IntegrationOutcome> outcomes;
    std::vector<bool> flipped(n_crystals);
    for (int c = 0; c < n_crystals; ++c) {
        flipped[c] = rng() % 2;
        UnitCell u = Scattered(cell, 0.01, rng);
        if (std::fabs(u.a - cell.b) + std::fabs(u.b - cell.a) < std::fabs(u.a - cell.a) + std::fabs(u.b - cell.b))
            std::swap(u.a, u.b);
        outcomes.push_back(Crystal(structure, key, cell, laws[0], flipped[c], CrystalLattice(u), rng));
    }

    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 4);
    CHECK(result.crystals == n_crystals);
    int agree = 0;
    for (int c = 0; c < n_crystals; ++c)
        agree += (result.choice[c] != 0) == flipped[c];
    CHECK(std::max(agree, n_crystals - agree) == n_crystals);
    CHECK(result.decided_fraction > 0.9);
}

TEST_CASE("Reindex: with a reference, cells that cannot tell the axes apart leave the indexing to the intensities", "[reindex]") {
    // The crystals of the test above, with a reference written in their true indexing: every crystal ends in
    // it - the reference fixes the run's indexing as a whole - and the cells, which only echo the indexer's
    // choice, put none of them elsewhere.
    const UnitCell cell{78.0f, 78.16f, 37.0f, 90.0f, 90.0f, 90.0f};
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(16);   // P222
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 1);

    const HKLKeyGenerator key(true, x.GetSpaceGroupOrP1());
    std::mt19937 rng(19);
    const Structure structure = MakeStructure(key, rng);
    std::vector<MergedReflection> reference;
    std::unordered_set<uint64_t> seen;
    for (const auto &t : structure.hkls) {
        if (!seen.insert(key(t[0], t[1], t[2]).pack()).second)
            continue;
        MergedReflection r;
        r.h = t[0]; r.k = t[1]; r.l = t[2];
        r.I = static_cast<float>(structure.intensity.at(key(t[0], t[1], t[2]).pack()));
        r.sigma = 1.0f;
        r.d = 1.0f / std::sqrt(t[0] * t[0] / (cell.a * cell.a) + t[1] * t[1] / (cell.b * cell.b)
                               + t[2] * t[2] / (cell.c * cell.c));
        reference.push_back(r);
    }

    const int n_crystals = 200;
    std::vector<IntegrationOutcome> outcomes;
    std::vector<bool> flipped(n_crystals);
    for (int c = 0; c < n_crystals; ++c) {
        flipped[c] = rng() % 10 < 7;   // most of them: the indexing that moves the fewest is the wrong one
        UnitCell u = Scattered(cell, 0.01, rng);
        if (std::fabs(u.a - cell.b) + std::fabs(u.b - cell.a) < std::fabs(u.a - cell.a) + std::fabs(u.b - cell.b))
            std::swap(u.a, u.b);
        outcomes.push_back(Crystal(structure, key, cell, laws[0], flipped[c], CrystalLattice(u), rng));
    }
    const auto as_indexed = outcomes;
    const auto moved = [&](int c) {
        const auto &r = outcomes[c].reflections[0], &r0 = as_indexed[c].reflections[0];
        return r.h != r0.h || r.k != r0.k || r.l != r0.l;
    };

    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, reference, 4);
    CHECK(result.crystals == n_crystals);
    CHECK(result.reference_overruled == 0);
    for (int c = 0; c < n_crystals; ++c)
        CHECK(moved(c) == flipped[c]);
}

TEST_CASE("Reindex: a monoclinic crystal on a pseudo-tetragonal lattice is relabelled as a whole only by a symmetry of its merge",
          "[reindex]") {
    // P21 with a and b equal: four indexings. Relabelling a whole merge by a diagonal operator would put the
    // two-fold axis on a, so it is no relabelling of a P21 merge - taken as one, it splits crystals that differ
    // by the two-fold into different indexings. Most crystals are truly in a diagonal indexing here, so the
    // indexing that moves the fewest of them is such an operator.
    const UnitCell cell{67.0f, 67.0f, 71.0f, 90.0f, 90.0f, 90.0f};
    DiffractionExperiment x;
    x.SetUnitCell(cell).SpaceGroupNumber(4);   // P21
    const auto laws = ReindexAmbiguityOperators(cell, x.GetSpaceGroupOrP1());
    REQUIRE(laws.size() == 3);
    std::vector<gemmi::Op> indexings{gemmi::Op::identity()};
    indexings.insert(indexings.end(), laws.begin(), laws.end());

    const HKLKeyGenerator key(true, x.GetSpaceGroupOrP1());
    std::mt19937 rng(23);
    std::exponential_distribution<double> wilson(1.0);
    std::normal_distribution<double> noise(0.0, 0.3);
    std::unordered_map<uint64_t, double> structure;
    std::vector<std::array<int, 3>> hkls;
    for (int h = -12; h <= 12; ++h)
        for (int k = -12; k <= 12; ++k)
            for (int l = -6; l <= 6; ++l)
                if (h != 0 || k != 0 || l != 0) {
                    hkls.push_back({h, k, l});
                    structure.try_emplace(key(h, k, l).pack(), wilson(rng));
                }

    const int n_crystals = 300;
    std::vector<IntegrationOutcome> outcomes(n_crystals);
    std::vector<int> truth(n_crystals);
    for (int c = 0; c < n_crystals; ++c) {
        truth[c] = c % 10 < 7 ? 2 : static_cast<int>(rng() % 4);
        outcomes[c].latt = CrystalLattice(cell);
        for (int i = 0; i < 150; ++i) {
            const auto &t = hkls[rng() % hkls.size()];
            const auto m = indexings[truth[c]].apply_to_hkl({{t[0], t[1], t[2]}});
            Reflection r{};
            r.h = m[0]; r.k = m[1]; r.l = m[2];
            r.d = 1.0f / std::sqrt((t[0] * t[0] + t[1] * t[1]) / (cell.a * cell.a) + t[2] * t[2] / (cell.c * cell.c));
            r.I = static_cast<float>(structure.at(key(t[0], t[1], t[2]).pack()) * (1.0 + noise(rng)));
            r.sigma = 0.1f;
            r.partiality = 1.0f;
            r.prescaling_corr = 1.0f;
            outcomes[c].reflections.push_back(r);
        }
    }

    const auto result = ResolveIndexingAmbiguity(outcomes, x, cell, laws, {}, 4);
    // Consistent: a general reflection, put into each crystal's indexing as indexed and then as chosen, lands
    // on the same reflection of the merge for every crystal.
    const std::array<int, 3> general{2, 5, 11};
    const auto merged = [&](int c) {
        const auto m = indexings[result.choice[c]].apply_to_hkl(indexings[truth[c]].apply_to_hkl(general));
        return key(m[0], m[1], m[2]).pack();
    };
    int consistent = 0;
    for (int c = 0; c < n_crystals; ++c)
        consistent += merged(c) == merged(0);
    CHECK(consistent == n_crystals);
}