// 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 <cmath>
#include <map>
#include <random>

#include "../rugnux/RugnuxJointSweeps.h"

namespace {

// A tetragonal P4 crystal: Laue class 4/m on a 4/mmm lattice, so (k, h, -l) indexes it a second way.
const gemmi::SpaceGroup &P4() { return *gemmi::find_spacegroup_by_name("P 4"); }
constexpr float A = 60.0f, C = 40.0f;

std::array<int, 3> P4Key(int h, int k, int l) {
    // the 4/m ASU representative, enough to give equivalents one intensity
    std::array<int, 3> best{h, k, l};
    const std::array<std::array<int, 3>, 8> eq{{{h, k, l}, {-k, h, l}, {-h, -k, l}, {k, -h, l},
                                                {-h, -k, -l}, {k, -h, -l}, {h, k, -l}, {-k, h, -l}}};
    for (const auto &e: eq) best = std::max(best, e);
    return best;
}

// Four observations of every reflection, with counting noise of `sigma`.
JointMemberSweep MakeSweep(const std::map<std::array<int, 3>, double> &truth, const gemmi::Op &indexed_as,
                           int n_frames, uint32_t seed, double sigma = 10.0) {
    std::mt19937 rng(seed);
    std::normal_distribution<double> noise(0.0, sigma);
    JointMemberSweep s;
    s.n_frames = n_frames;
    s.cell = UnitCell{A, A, C, 90, 90, 90};
    s.space_group = P4();
    for (const auto &[hkl, I]: truth) {
        const auto m = indexed_as.apply_to_hkl({{hkl[0], hkl[1], hkl[2]}});
        const auto d = static_cast<float>(1.0 / std::sqrt((hkl[0] * hkl[0] + hkl[1] * hkl[1]) / (A * A)
                                                          + hkl[2] * hkl[2] / (C * C)));
        double sum = 0;
        for (int i = 0; i < 4; ++i) {
            ScaledFull f;
            f.h = m[0]; f.k = m[1]; f.l = m[2];
            f.I = static_cast<float>(I + noise(rng));
            f.sigma = static_cast<float>(sigma);
            f.image_number = static_cast<float>(rng() % n_frames);
            f.var_bkg = 4.0f;
            f.var_per_I = 0.5f;
            f.capture = 0.0f;
            f.d = d;
            s.fulls.push_back(f);
            sum += f.I;
        }
        MergedReflection r;
        r.h = m[0]; r.k = m[1]; r.l = m[2];
        r.I = static_cast<float>(sum / 4);
        r.sigma = static_cast<float>(sigma / 2);
        r.d = d;
        s.merged.push_back(r);
    }
    return s;
}

std::map<std::array<int, 3>, double> Truth(uint32_t seed) {
    std::mt19937 rng(seed);
    std::exponential_distribution<double> wilson(1.0);
    std::map<std::array<int, 3>, double> t;
    for (int h = -12; h <= 12; ++h)
        for (int k = -12; k <= 12; ++k)
            for (int l = -8; l <= 8; ++l) {
                if (!(h || k || l)) continue;
                const auto key = P4Key(h, k, l);
                if (!t.contains(key)) t[key] = 1000.0 * wilson(rng);
            }
    return t;
}

}  // namespace

TEST_CASE("RelateJointSweeps_AmbiguityRecoveredAndAStrangerLeftOut", "[Rugnux]") {
    const auto truth = Truth(1);
    const gemmi::Op twin = gemmi::parse_triplet("y,x,-z");   // h,k,l -> k,h,-l
    std::vector<JointMemberSweep> sweeps;
    sweeps.push_back(MakeSweep(truth, gemmi::Op::identity(), 100, 2));
    sweeps.push_back(MakeSweep(truth, twin, 80, 3));
    sweeps.push_back(MakeSweep(Truth(99), gemmi::Op::identity(), 50, 4));   // another crystal
    sweeps[0].isa = 20.0;                                                   // the reference: the best ISa
    sweeps[1].isa = 10.0;
    sweeps[2].isa = 15.0;

    const auto relation = RelateJointSweeps(sweeps);
    REQUIRE(relation.reference == 0);
    CHECK(relation.sweeps[0].included);
    REQUIRE(relation.sweeps[1].included);
    CHECK(relation.sweeps[1].cc > 0.99);
    CHECK(relation.sweeps[1].identity_cc < 0.5);
    // the operator found takes sweep 2's indices back onto the reference's (up to the Laue group)
    const auto back = relation.sweeps[1].cob.apply_to_hkl(twin.apply_to_hkl({{3, 1, 2}}));
    CHECK(P4Key(back[0], back[1], back[2]) == P4Key(3, 1, 2));
    CHECK(relation.sweeps[1].decision == "isomorphous");
    CHECK_FALSE(relation.sweeps[2].included);
    CHECK(relation.sweeps[2].decision == "not isomorphous");

    const auto input = BuildJointMergeInput(sweeps, relation);
    REQUIRE(input.segments.size() == 2);
    CHECK(input.segments[0].first_frame == 0);
    CHECK(input.segments[0].n_frames == 100);
    CHECK(input.segments[1].first_frame == 100);
    CHECK(input.segments[1].n_frames == 80);
    CHECK(input.segment_sweep == std::vector<int>{0, 1});
    REQUIRE(input.outcomes.size() == 180);
    // a full as the combine takes it: count I/c at correction c = var_per_I, variance var_bkg / c^2
    const Reflection *r = nullptr;
    for (const auto &o: input.outcomes)
        if (!o.reflections.empty()) { r = &o.reflections.front(); break; }
    REQUIRE(r != nullptr);
    CHECK(r->image_scale_corr == Catch::Approx(0.5));
    CHECK(r->var_bkg == Catch::Approx(16.0));
    CHECK(r->partiality == 1.0f);
}

TEST_CASE("RelateJointSweeps_DoubtfulSweepDecidedOnTheMerge", "[Rugnux]") {
    // Sweep 1: the same crystal measured noisily, its ratio uncertain; sweep 2: precise data of a crystal
    // whose intensities differ by a few percent from the reference's. Both fall in the doubtful band, and
    // the held-out residual keeps the first and leaves out the second.
    const auto truth = Truth(1);
    auto shifted = truth;
    std::mt19937 rng(7);
    std::normal_distribution<double> change(0.0, 0.14);
    for (auto &[hkl, I]: shifted)
        I *= std::exp(change(rng));
    std::vector<JointMemberSweep> sweeps;
    sweeps.push_back(MakeSweep(truth, gemmi::Op::identity(), 100, 2, 100.0));
    sweeps.push_back(MakeSweep(truth, gemmi::Op::identity(), 100, 3, 1500.0));
    sweeps.push_back(MakeSweep(shifted, gemmi::Op::identity(), 100, 4, 10.0));
    sweeps[0].isa = 20.0;

    const auto relation = RelateJointSweeps(sweeps);
    REQUIRE(relation.reference == 0);
    INFO("ratios " << relation.sweeps[1].iso_ratio << " +- " << relation.sweeps[1].iso_se << ", "
         << relation.sweeps[2].iso_ratio << " +- " << relation.sweeps[2].iso_se);
    CHECK(relation.sweeps[1].decision == "doubtful: kept");
    CHECK(relation.sweeps[1].included);
    CHECK(relation.sweeps[1].held_out_change < 0.0);
    CHECK(relation.sweeps[2].decision == "doubtful: left out");
    CHECK_FALSE(relation.sweeps[2].included);
    CHECK(relation.sweeps[2].held_out_change > 0.0);
}