// 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 "../common/CrystalLattice.h"
#include "../common/Coord.h"
#include "../common/UnitCell.h"
#include "../image_analysis/lattice_search/LePageLattice.h"
#include "gemmi/symmetry.hpp"
#include <array>
#include <cmath>
#include <random>

namespace {
int CentringMultiplicity(char c) {
    switch (c) {
        case 'A': case 'B': case 'C': case 'I': return 2;
        case 'R': return 3;
        case 'F': return 4;
        default:  return 1;
    }
}

// The lattice is what has to come back, not the axes: a conventional cell of the right class whose
// PRIMITIVE volume is the one we started from. Comparing lengths element-wise would fail a correct
// answer given in another setting.
void CheckLattice(const std::optional<LePageResult> &r, gemmi::CrystalSystem system, char centering,
                  float primitive_volume) {
    REQUIRE(r.has_value());
    CHECK(r->system == system);
    CHECK(r->centering == centering);
    const float v = std::fabs(r->conventional.CalcVolume()) / CentringMultiplicity(r->centering);
    CHECK(v == Catch::Approx(primitive_volume).epsilon(0.02));
}

// The Miller-index matrix a filter is handed, against the integer matrix expected there.
bool SameMatrix(const gemmi::Mat33 &m, const std::array<int, 9> &expected) {
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++)
            if (std::fabs(m[i][j] - expected[3 * i + j]) > 1e-6)
                return false;
    return true;
}

// The lattice as an indexer hands it over: some basis of it, in some orientation, with noise.
CrystalLattice Present(const CrystalLattice &L, std::mt19937 &rng, float noise_A) {
    std::uniform_int_distribution<int> pick(0, 2), amount(-1, 1);
    std::uniform_real_distribution<float> uni(0, 1);
    std::normal_distribution<float> gauss(0, noise_A);
    gemmi::Mat33 m(1, 0, 0, 0, 1, 0, 0, 0, 1);
    for (int n = 0; n < 4; n++) {
        const int i = pick(rng), j = pick(rng);
        if (i == j)
            continue;
        gemmi::Mat33 shear(1, 0, 0, 0, 1, 0, 0, 0, 1);
        shear.a[i][j] = amount(rng);
        m = shear.multiply(m);
    }
    const float theta = 2 * (float)M_PI * uni(rng), phi = std::acos(2 * uni(rng) - 1);
    const Coord axis(std::sin(phi) * std::cos(theta), std::sin(phi) * std::sin(theta), std::cos(phi));
    CrystalLattice out = L.Multiply(m).Multiply(RotMatrix(2 * (float)M_PI * uni(rng), axis));
    Coord v[3] = {out.Vec0(), out.Vec1(), out.Vec2()};
    for (auto &k : v) { k.x += gauss(rng); k.y += gauss(rng); k.z += gauss(rng); }
    return CrystalLattice(v[0], v[1], v[2]);
}
}   // namespace

TEST_CASE("LePageLattice - the fourteen Bravais lattices") {
    struct Case { gemmi::CrystalSystem system; char centering; float a, b, c, al, be, ga; };
    const Case cases[] = {
        {gemmi::CrystalSystem::Triclinic,    'P', 23, 31, 41, 81, 95, 71},
        {gemmi::CrystalSystem::Monoclinic,   'P', 31, 43, 57, 90, 103, 90},
        {gemmi::CrystalSystem::Monoclinic,   'C', 91, 43, 57, 90, 103, 90},
        {gemmi::CrystalSystem::Orthorhombic, 'P', 31, 43, 57, 90, 90, 90},
        {gemmi::CrystalSystem::Orthorhombic, 'C', 31, 43, 57, 90, 90, 90},
        {gemmi::CrystalSystem::Orthorhombic, 'I', 31, 43, 57, 90, 90, 90},
        {gemmi::CrystalSystem::Orthorhombic, 'F', 31, 43, 57, 90, 90, 90},
        {gemmi::CrystalSystem::Tetragonal,   'P', 47, 47, 71, 90, 90, 90},
        {gemmi::CrystalSystem::Tetragonal,   'I', 47, 47, 71, 90, 90, 90},
        {gemmi::CrystalSystem::Trigonal,     'R', 61, 61, 133, 90, 90, 120},
        {gemmi::CrystalSystem::Hexagonal,    'P', 61, 61, 97, 90, 90, 120},
        {gemmi::CrystalSystem::Cubic,        'P', 71, 71, 71, 90, 90, 90},
        {gemmi::CrystalSystem::Cubic,        'I', 71, 71, 71, 90, 90, 90},
        {gemmi::CrystalSystem::Cubic,        'F', 71, 71, 71, 90, 90, 90},
    };
    for (const Case &c : cases) {
        const CrystalLattice conventional(c.a, c.b, c.c, c.al, c.be, c.ga);
        const CrystalLattice primitive = conventional.ToPrimitive(c.centering);
        const float primitive_volume = std::fabs(primitive.CalcVolume());
        std::mt19937 rng(20260831);
        for (int i = 0; i < 20; i++) {
            INFO("class " << (int)c.system << c.centering << " presentation " << i);
            CheckLattice(LePageLattice(Present(primitive, rng, 0.02f)), c.system, c.centering,
                         primitive_volume);
        }
    }
}

TEST_CASE("LePageLattice - a cubic F lattice on the Niggli type boundary") {
    // An fcc lattice has both a 60/60/60 and a ~120/90/120 shortest-vector basis, so it sits ON the
    // boundary between the two Niggli types by construction and the reduction lands on either side
    // according to the last bits of the cell it is given. Reading the symmetry off the metric has no
    // forms to fall between, so the answer does not depend on which side it landed on.
    const CrystalLattice conventional(121.0f * std::sqrt(2.0f), 121.0f * std::sqrt(2.0f),
                                      121.0f * std::sqrt(2.0f), 90, 90, 90);
    const CrystalLattice primitive = conventional.ToPrimitive('F');
    const float primitive_volume = std::fabs(primitive.CalcVolume());
    std::mt19937 rng(7);
    for (int i = 0; i < 40; i++) {
        INFO("presentation " << i);
        CheckLattice(LePageLattice(Present(primitive, rng, 0.1f)), gemmi::CrystalSystem::Cubic, 'F',
                     primitive_volume);
    }
}

TEST_CASE("LePageLattice - a tetragonal I description of a cubic F lattice") {
    // Same lattice as above, handed over in the setting a, a, a*sqrt(2) that describes it as body-
    // centred tetragonal. Both descriptions are the same lattice, and the answer has to be the same.
    const float a = 120.5f;
    const CrystalLattice tetragonal(a, a, a * std::sqrt(2.0f), 90, 90, 90);
    const CrystalLattice primitive = tetragonal.ToPrimitive('I');
    auto r = LePageLattice(primitive);
    REQUIRE(r.has_value());
    CHECK(r->system == gemmi::CrystalSystem::Cubic);
    CHECK(r->centering == 'F');
    CHECK(std::fabs(r->conventional.CalcVolume()) / 4 ==
          Catch::Approx(std::fabs(primitive.CalcVolume())).epsilon(0.01));
}

TEST_CASE("LePageLattice - a monoclinic lattice is named C even where I is less oblique") {
    // The least oblique naming of this lattice's (a, c) plane is body-centred, beta a few degrees
    // from 90. The answer must still be the C-centred reference setting - the space-group search
    // enumerates reference settings only, so a promotion earned by a lattice named I can never be
    // confirmed. And it must be the least oblique C the plane offers, not a relabel: on this
    // lattice a -> a + c puts beta past the 150-degree bound of the constrained refinement.
    const CrystalLattice conventional(40, 54, 76, 90, 93, 90);   // the I naming of the lattice
    const CrystalLattice primitive = conventional.ToPrimitive('I');
    const float primitive_volume = std::fabs(primitive.CalcVolume());
    std::mt19937 rng(20260908);
    for (int i = 0; i < 20; i++) {
        INFO("presentation " << i);
        const auto r = LePageLattice(Present(primitive, rng, 0.02f));
        CheckLattice(r, gemmi::CrystalSystem::Monoclinic, 'C', primitive_volume);
        const UnitCell uc = r->conventional.GetUnitCell();
        CHECK(uc.beta > 90.0f);
        CHECK(uc.beta < 150.0f);
    }
}

TEST_CASE("LePageLattice - the change of basis is integral and right-handed") {
    const CrystalLattice conventional(47, 47, 71, 90, 90, 90);
    const CrystalLattice primitive = conventional.ToPrimitive('I');
    auto r = LePageLattice(primitive);
    REQUIRE(r.has_value());
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++)
            CHECK(r->reindex[i][j] == Catch::Approx(std::round(r->reindex[i][j])).margin(1e-9));
    CHECK(r->reindex.determinant() > 0);
    CHECK(r->conventional.CalcVolume() > 0);
}

TEST_CASE("LePageLattice - a pseudo-symmetric metric is not promoted") {
    // A monoclinic cell whose beta sits a few degrees from 90 is not orthorhombic, however close the
    // reduced form is to an orthorhombic character.
    const CrystalLattice L(31, 43, 57, 90, 93, 90);
    auto r = LePageLattice(L);
    REQUIRE(r.has_value());
    CHECK(r->system == gemmi::CrystalSystem::Monoclinic);
    CHECK(r->centering == 'P');
}

TEST_CASE("LePageLattice - an operator filter picks out a sub-lattice of the metric") {
    // A hexagonal metric carries twelve rotations, and asked as it stands that is the answer. Keeping
    // only the three two-folds that close into 222 - the one along c and the two in the plane, along
    // a+b and a-b - leaves the orthorhombic sub-lattice, whose conventional cell is (a+b, a-b, c) on
    // twice the volume and so C-centred.
    const CrystalLattice L(100, 100, 70, 90, 90, 120);
    const auto whole_metric = LePageLattice(L);
    REQUIRE(whole_metric.has_value());
    CHECK(whole_metric->system == gemmi::CrystalSystem::Hexagonal);

    const std::array<int, 9> along_c = {-1, 0, 0, 0, -1, 0, 0, 0, 1};       // -h,-k,l
    const std::array<int, 9> along_a_plus_b = {0, 1, 0, 1, 0, 0, 0, 0, -1};  // k,h,-l
    const std::array<int, 9> along_a_minus_b = {0, -1, 0, -1, 0, 0, 0, 0, -1};   // -k,-h,-l
    const auto r = LePageLattice(L, LATTICE_MAX_OBLIQUITY_DEG, [&](const gemmi::Mat33 &m) {
        return SameMatrix(m, along_c) || SameMatrix(m, along_a_plus_b) ||
               SameMatrix(m, along_a_minus_b);
    });
    REQUIRE(r.has_value());
    CHECK(r->system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(r->centering == 'C');
    const UnitCell uc = r->conventional.GetUnitCell();
    CHECK(uc.a == Catch::Approx(100).epsilon(0.01));
    CHECK(uc.b == Catch::Approx(100 * std::sqrt(3.0)).epsilon(0.01));
    CHECK(uc.c == Catch::Approx(70).epsilon(0.01));
}

TEST_CASE("LePageLattice - a filtered group may not rest on an operator the filter refused") {
    // Two two-folds sixty degrees apart generate a three-fold, so the pair closes into the whole
    // hexagonal group - four of whose two-folds the filter rejected. A sub-lattice whose own operators
    // were offered and refused is not one the filter supports, so there is no answer to give.
    const CrystalLattice L(100, 100, 70, 90, 90, 120);
    const std::array<int, 9> along_a_plus_b = {0, 1, 0, 1, 0, 0, 0, 0, -1};   // k,h,-l
    const std::array<int, 9> along_a = {1, 0, 0, -1, -1, 0, 0, 0, -1};        // h,-h-k,-l
    const auto r = LePageLattice(L, LATTICE_MAX_OBLIQUITY_DEG, [&](const gemmi::Mat33 &m) {
        return SameMatrix(m, along_a_plus_b) || SameMatrix(m, along_a);
    });
    CHECK(!r.has_value());
}