// 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/CrystalSetting.h"
#include "../image_analysis/scale_merge/HKLKey.h"
#include "../image_analysis/scale_merge/ReindexAmbiguity.h"
#include "../image_analysis/scale_merge/RfreeFlags.h"

namespace {
    const gemmi::SpaceGroup &Named(const char *name) { return *gemmi::find_spacegroup_by_name(name); }

    // The group and cell a change of basis writes, as the pipeline would write them.
    struct Written {
        std::string group;
        UnitCell cell;
    };
    Written Apply(const UnitCell &cell, const gemmi::SpaceGroup &sg, const gemmi::Op &cob) {
        const gemmi::SpaceGroup *moved = SpaceGroupInBasis(sg, cob);
        REQUIRE(moved != nullptr);
        return {moved->xhm(), CellInBasis(cell, cob)};
    }
}

TEST_CASE("Setting: P21212 with the pure axis shortest is written as P 21 21 2", "[setting]") {
    // The a<b<c cell puts the pure two-fold on a; XDS and ITA put it on c, with a<b.
    const UnitCell cell{52.51f, 87.87f, 137.72f, 90.0f, 90.0f, 90.0f};
    const auto w = Apply(cell, Named("P 2 21 21"), ChooseOutputSetting(cell, Named("P 2 21 21")));
    CHECK(w.group == "P 21 21 2");
    CHECK(w.cell.a == Catch::Approx(87.87));
    CHECK(w.cell.b == Catch::Approx(137.72));
    CHECK(w.cell.c == Catch::Approx(52.51));
}

TEST_CASE("Setting: P2221 with the screw on a moves it to c with a<b", "[setting]") {
    const UnitCell cell{40.0f, 50.0f, 60.0f, 90.0f, 90.0f, 90.0f};
    const auto w = Apply(cell, Named("P 21 2 2"), ChooseOutputSetting(cell, Named("P 21 2 2")));
    CHECK(w.group == "P 2 2 21");
    CHECK(w.cell.a == Catch::Approx(50.0));
    CHECK(w.cell.b == Catch::Approx(60.0));
    CHECK(w.cell.c == Catch::Approx(40.0));
}

TEST_CASE("Setting: groups with alike axes are ordered a<=b<=c, and a standard run is left alone", "[setting]") {
    const UnitCell unordered{121.0f, 199.0f, 189.0f, 90.0f, 90.0f, 90.0f};
    const auto w = Apply(unordered, Named("I 2 2 2"), ChooseOutputSetting(unordered, Named("I 2 2 2")));
    CHECK(w.group == "I 2 2 2");
    CHECK(w.cell.a == Catch::Approx(121.0));
    CHECK(w.cell.b == Catch::Approx(189.0));
    CHECK(w.cell.c == Catch::Approx(199.0));

    const UnitCell standard{50.0f, 60.0f, 70.0f, 90.0f, 90.0f, 90.0f};
    CHECK(ChooseOutputSetting(standard, Named("P 21 21 21")) == gemmi::Op::identity());
    const UnitCell tetragonal{78.0f, 78.0f, 37.0f, 90.0f, 90.0f, 90.0f};
    CHECK(ChooseOutputSetting(tetragonal, Named("P 43 21 2")) == gemmi::Op::identity());
}

TEST_CASE("Setting: C222 is written with a<b", "[setting]") {
    const UnitCell cell{222.0f, 131.0f, 86.0f, 90.0f, 90.0f, 90.0f};
    const auto w = Apply(cell, Named("C 2 2 2"), ChooseOutputSetting(cell, Named("C 2 2 2")));
    CHECK(w.group == "C 2 2 2");
    CHECK(w.cell.a == Catch::Approx(131.0));
    CHECK(w.cell.b == Catch::Approx(222.0));
}

TEST_CASE("Setting: monoclinic c-unique is written b-unique with beta >= 90", "[setting]") {
    const UnitCell cell{35.0f, 40.0f, 44.0f, 90.0f, 90.0f, 110.0f};
    const auto w = Apply(cell, Named("P 1 1 21"), ChooseOutputSetting(cell, Named("P 1 1 21")));
    CHECK(w.group == "P 1 21 1");
    CHECK(w.cell.b == Catch::Approx(44.0));
    CHECK(w.cell.beta == Catch::Approx(110.0).margin(1e-3));
}

TEST_CASE("Setting: C2 with an acute beta is made obtuse; a standard C2 is left alone", "[setting]") {
    const UnitCell acute{120.0f, 50.0f, 75.0f, 90.0f, 87.5f, 90.0f};
    const auto w = Apply(acute, Named("C 1 2 1"), ChooseOutputSetting(acute, Named("C 1 2 1")));
    CHECK(w.group == "C 1 2 1");
    CHECK(w.cell.beta == Catch::Approx(92.5).margin(1e-3));

    const UnitCell standard{112.7f, 52.8f, 44.5f, 90.0f, 103.1f, 90.0f};
    CHECK(ChooseOutputSetting(standard, Named("C 1 2 1")) == gemmi::Op::identity());
}

TEST_CASE("Setting: I2 is written as C2, no more oblique", "[setting]") {
    const UnitCell c2{112.7f, 52.8f, 44.5f, 90.0f, 103.1f, 90.0f};
    // The same lattice described I-centred.
    gemmi::Op to_i = gemmi::Op::identity();
    const gemmi::SpaceGroup *i2 = nullptr;
    for (const gemmi::Op &cob : UnimodularOperators()) {
        gemmi::GroupOps g = Named("C 1 2 1").operations();
        g.change_basis_forward(cob);
        if (g.find_centering() != 'I')
            continue;
        if (const gemmi::SpaceGroup *m = SpaceGroupInBasis(Named("C 1 2 1"), cob); m && m->xhm() == "I 1 2 1") {
            to_i = cob;
            i2 = m;
            break;
        }
    }
    REQUIRE(i2 != nullptr);
    const UnitCell i_cell = CellInBasis(c2, to_i);
    const auto w = Apply(i_cell, *i2, ChooseOutputSetting(i_cell, *i2));
    CHECK(w.group == "C 1 2 1");
    CHECK(w.cell.beta >= 90.0f);
    CHECK(w.cell.beta <= 103.1f + 1e-3f);
    CHECK(w.cell.b == Catch::Approx(52.8));
}

TEST_CASE("Setting: triclinic keeps the reduced cell", "[setting]") {
    const UnitCell cell{27.3f, 31.9f, 34.3f, 88.0f, 71.6f, 68.0f};
    CHECK(ChooseOutputSetting(cell, gemmi::get_spacegroup_p1()) == gemmi::Op::identity());
}

TEST_CASE("Setting: a cell given with -C sets the axis order, and a non-standard -S keeps its setting", "[setting]") {
    const UnitCell cell{87.87f, 137.72f, 52.51f, 90.0f, 90.0f, 90.0f};   // standard P 21 21 2
    const UnitCell given{52.5f, 87.9f, 137.7f, 90.0f, 90.0f, 90.0f};
    const auto w = Apply(cell, Named("P 21 21 2"), ChooseOutputSetting(cell, Named("P 21 21 2"), nullptr, given));
    CHECK(w.group == "P 2 21 21");
    CHECK(w.cell.a == Catch::Approx(52.51));
    CHECK(w.cell.b == Catch::Approx(87.87));

    // -S P 2 21 21 without a cell: the pure axis on a, the other two ordered.
    const auto s = Apply(cell, Named("P 21 21 2"),
                         ChooseOutputSetting(cell, Named("P 21 21 2"), &Named("P 2 21 21")));
    CHECK(s.group == "P 2 21 21");
    CHECK(s.cell.a == Catch::Approx(52.51));
    CHECK(s.cell.b == Catch::Approx(87.87));
}

TEST_CASE("Setting: a fixed P21212 is seated on the axes its absences name", "[setting]") {
    // a<b<c, pure two-fold on a: h00 odd present, 0k0 and 00l odd absent.
    const UnitCell cell{52.0f, 88.0f, 138.0f, 90.0f, 90.0f, 90.0f};
    std::vector<ReflectionZ> z;
    for (int n = 1; n <= 12; n++) {
        z.push_back({{n, 0, 0}, 20.0});
        z.push_back({{0, n, 0}, n % 2 ? 0.1 : 20.0});
        z.push_back({{0, 0, n}, n % 2 ? -0.2 : 20.0});
        z.push_back({{n, n, 0}, 15.0});   // a zone reflection no candidate predicts absent
    }
    const gemmi::Op cob = SeatGroupByAbsences(cell, Named("P 21 21 2"), z);
    // P 21 21 2 as given describes the data once they are put through cob: the pure axis is now c.
    CHECK(CellInBasis(cell, cob).c == Catch::Approx(52.0));

    // P2221 whose screw row (00l) was measured absent while 0k0 was never recorded: the screw is on c.
    std::vector<ReflectionZ> z2;
    for (int n = 1; n <= 12; n++) {
        z2.push_back({{n, 0, 0}, 20.0});
        z2.push_back({{0, 0, n}, n % 2 ? 0.2 : 20.0});
        z2.push_back({{n, n, 0}, 15.0});
    }
    CHECK(SeatGroupByAbsences(cell, Named("P 2 2 21"), z2) == gemmi::Op::identity());
}

TEST_CASE("Setting: a reference on permuted axes is matched, and its free flags follow the right reflections",
          "[setting][reindex]") {
    // The reference: P 21 21 2 on 88/138/52, one distinct intensity and a flag per reflection.
    const UnitCell ref_cell{88.0f, 138.0f, 52.0f, 90.0f, 90.0f, 90.0f};
    const gemmi::SpaceGroup &ref_sg = Named("P 21 21 2");
    const HKLKeyGenerator key(true, ref_sg);
    const gemmi::UnitCell g(ref_cell);
    std::vector<MergedReflection> reference;
    std::unordered_set<uint64_t> seen;
    for (int h = 0; h <= 10; h++)
        for (int k = 0; k <= 10; k++)
            for (int l = 0; l <= 8; l++) {
                if (h + k + l == 0 || ref_sg.operations().is_systematically_absent({h, k, l}))
                    continue;
                const uint64_t kk = key(h, k, l).pack();
                if (!seen.insert(kk).second)
                    continue;
                MergedReflection r{};
                r.h = h; r.k = k; r.l = l;
                r.I = static_cast<float>(100 + (kk * 2654435761u) % 9973);
                r.sigma = 1.0f;
                r.d = static_cast<float>(g.calculate_d({h, k, l}));
                r.rfree_flag = (kk % 7) == 0;
                reference.push_back(r);
            }

    // The data: the same reflections on a<b<c axes (a = the reference's c), as rugnux indexes them.
    const UnitCell data_cell{52.0f, 88.0f, 138.0f, 90.0f, 90.0f, 90.0f};
    std::vector<MergedReflection> data = reference;
    for (auto &r : data) {
        const int h = r.h, k = r.k, l = r.l;
        r.h = l; r.k = h; r.l = k;
        r.rfree_flag = false;
    }

    // Matched as indexed, the flags would land on unrelated reflections - and the match says so.
    CHECK(MatchReference(data, reference, ref_sg).cc < MIN_REFERENCE_CC);

    const auto choice = ChooseReferenceSetting(data, data_cell, Named("P 2 21 21"), reference, ref_cell, ref_sg);
    CHECK(choice.cc == Catch::Approx(1.0).margin(1e-6));
    auto moved = ReindexMergedIntoAsu(data, HklOperator(choice.cob), ref_sg, true);
    const auto match = MatchReference(moved, reference, ref_sg);
    CHECK(match.cc >= MIN_REFERENCE_CC);
    CHECK(match.matched_fraction == Catch::Approx(1.0));
    ApplyReferenceFreeFlags(moved, ref_sg, reference);
    for (size_t i = 0; i < moved.size(); i++)   // ReindexMergedIntoAsu keeps the order
        CHECK(moved[i].rfree_flag == reference[i].rfree_flag);
}