// SPDX-FileCopyrightText: 2025 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/LatticeSearch.h"
#include "gemmi/symmetry.hpp"
#include <cmath>

// Helper: check near-equality of unit cell parameters
static void check_uc(const UnitCell& uc, double a, double b, double c,
                     double alpha, double beta, double gamma,
                     double eps_len = 1e-6, double eps_ang = 1e-4) {
    CHECK(uc.a == Catch::Approx(a).margin(eps_len));
    CHECK(uc.b == Catch::Approx(b).margin(eps_len));
    CHECK(uc.c == Catch::Approx(c).margin(eps_len));
    CHECK(uc.alpha == Catch::Approx(alpha).margin(eps_ang));
    CHECK(uc.beta  == Catch::Approx(beta ).margin(eps_ang));
    CHECK(uc.gamma == Catch::Approx(gamma).margin(eps_ang));
}


TEST_CASE("LatticeSearch - cubic I") {
    // Build a body-centered cubic cell with a=40:
    // primitive basis vectors (conventional I cubic primitive):
    // p1 = (0, a/2, a/2), p2 = (a/2, 0, a/2), p3 = (a/2, a/2, 0)
    const double a = 40.0;
    CrystalLattice L(
        Coord(a, 0, 0),
        Coord(0, a, 0),
        Coord(0, 0, a)
    );
    L = L.ToPrimitive('I');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Cubic);
    CHECK(res.centering == 'I');

    // Conventional cubic I should have equal edges and 90° angles
    auto uc = res.conventional.GetUnitCell();
    CHECK(uc.a == Catch::Approx( a )); // In this construction, conventional a matches given a
    CHECK(uc.b == Catch::Approx( a ));
    CHECK(uc.c == Catch::Approx( a ));
    CHECK(uc.alpha == Catch::Approx(90.0));
    CHECK(uc.beta  == Catch::Approx(90.0));
    CHECK(uc.gamma == Catch::Approx(90.0));
}

TEST_CASE("LatticeSearch - cubic F") {
    // Build a body-centered cubic cell with a=40:
    // primitive basis vectors (conventional I cubic primitive):
    // p1 = (0, a/2, a/2), p2 = (a/2, 0, a/2), p3 = (a/2, a/2, 0)
    const double a = 40.0;
    CrystalLattice L(
        Coord(a, 0, 0),
        Coord(0, a, 0),
        Coord(0, 0, a)
    );
    L = L.ToPrimitive('F');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Cubic);
    CHECK(res.centering == 'F');

    // Conventional cubic I should have equal edges and 90° angles
    auto uc = res.conventional.GetUnitCell();
    CHECK(uc.a == Catch::Approx( a )); // In this construction, conventional a matches given a
    CHECK(uc.b == Catch::Approx( a ));
    CHECK(uc.c == Catch::Approx( a ));
    CHECK(uc.alpha == Catch::Approx(90.0));
    CHECK(uc.beta  == Catch::Approx(90.0));
    CHECK(uc.gamma == Catch::Approx(90.0));
}

TEST_CASE("LatticeSearch - cubic P") {
    // Simple cubic P, a=30
    const double a = 30.0;
    CrystalLattice L(
        Coord(a,0,0),
        Coord(0,a,0),
        Coord(0,0,a)
    );

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Cubic);
    CHECK(res.centering == 'P');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, a, a, a, 90.0, 90.0, 90.0, 1e-6, 1e-4);
}

TEST_CASE("LatticeSearch - tetragonal I") {
    // Build a body-centered cubic cell with a=40:
    // primitive basis vectors (conventional I cubic primitive):
    // p1 = (0, a/2, a/2), p2 = (a/2, 0, a/2), p3 = (a/2, a/2, 0)
    const double a = 40.0;
    const double b = 34.0;
    CrystalLattice L(
        Coord(a, 0, 0),
        Coord(0, a, 0),
        Coord(0, 0, b)
    );
    L = L.ToPrimitive('I');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Tetragonal);
    CHECK(res.centering == 'I');

    // Conventional cubic I should have equal edges and 90° angles
    auto uc = res.conventional.GetUnitCell();
    CHECK(uc.a == Catch::Approx( a )); // In this construction, conventional a matches given a
    CHECK(uc.b == Catch::Approx( a ));
    CHECK(uc.c == Catch::Approx( b ));
    CHECK(uc.alpha == Catch::Approx(90.0));
    CHECK(uc.beta  == Catch::Approx(90.0));
    CHECK(uc.gamma == Catch::Approx(90.0));
}

TEST_CASE("LatticeSearch - tetragonal I - v2") {
    // Build a body-centered cubic cell with a=40:
    // primitive basis vectors (conventional I cubic primitive):
    // p1 = (0, a/2, a/2), p2 = (a/2, 0, a/2), p3 = (a/2, a/2, 0)
    const double a = 40.0;
    const double b = 54.0;
    CrystalLattice L(
        Coord(a, 0, 0),
        Coord(0, a, 0),
        Coord(0, 0, b)
    );
    L = L.ToPrimitive('I');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Tetragonal);
    CHECK(res.centering == 'I');

    // Conventional cubic I should have equal edges and 90° angles
    auto uc = res.conventional.GetUnitCell();
    CHECK(uc.a == Catch::Approx( a )); // In this construction, conventional a matches given a
    CHECK(uc.b == Catch::Approx( a ));
    CHECK(uc.c == Catch::Approx( b ));
    CHECK(uc.alpha == Catch::Approx(90.0));
    CHECK(uc.beta  == Catch::Approx(90.0));
    CHECK(uc.gamma == Catch::Approx(90.0));
}

// Tetragonal P: a=b!=c, all angles 90, P-centering
TEST_CASE("LatticeSearch - tetragonal P") {
    const double a = 37.0, c = 59.0;
    CrystalLattice L(
        Coord(a,0,0),
        Coord(0,a,0),
        Coord(0,0,c)
    );

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Tetragonal);
    CHECK(res.centering == 'P');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, a, a, c, 90.0, 90.0, 90.0, 1e-2, 1e-2);
}

// Orthorhombic F: all angles 90, unequal edges, F-centering
TEST_CASE("LatticeSearch - orthorhombic F") {
    const double a = 35.0, b = 41.0, c = 57.0;
    CrystalLattice conv(a,b,c, 90.0,90.0,90.0);

    CrystalLattice L = conv.ToPrimitive('F');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'F');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, a, b, c, 90.0, 90.0, 90.0, 1e-1, 1e-2);
}

TEST_CASE("LatticeSearch - orthorhombic F - permutation 1") {
    const double a = 41.0, b = 57.0, c = 35.0;
    CrystalLattice conv(a,b,c, 90.0,90.0,90.0);

    CrystalLattice L = conv.ToPrimitive('F');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'F');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, c, a, b, 90.0, 90.0, 90.0, 1e-1, 1e-2);
}

// Orthorhombic C: all angles 90, unequal edges, C-centering
TEST_CASE("LatticeSearch - orthorhombic C") {
    const double a = 35.0, b = 41.0, c = 57.0;
    CrystalLattice conv(a,b,c, 90.0,90.0,90.0);
    CrystalLattice L = conv.ToPrimitive('C');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'C');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, a, b, c, 90.0, 90.0, 90.0, 1e-1, 1e-2);
}

TEST_CASE("LatticeSearch - orthorhombic I") {
    const double a = 35.0, b = 41.0, c = 57.0;
    CrystalLattice conv(a,b,c, 90.0,90.0,90.0);
    CrystalLattice L = conv.ToPrimitive('I');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'I');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, a, b, c, 90.0, 90.0, 90.0, 1e-2, 1e-2);
}


TEST_CASE("LatticeSearch - orthorhombic I - permutation1") {
    const double a = 57.0, b = 41.0, c = 35.0;
    CrystalLattice conv(a,b,c, 90.0,90.0,90.0);
    CrystalLattice L = conv.ToPrimitive('I');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'I');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, c, b, a, 90.0, 90.0, 90.0, 1e-2, 1e-2);
}

TEST_CASE("LatticeSearch - orthorhombic I - permutation2") {
    const double a = 41.0, b = 57.0, c = 35.0;
    CrystalLattice conv(a,b,c, 90.0,90.0,90.0);
    CrystalLattice L = conv.ToPrimitive('I');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'I');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, c, a, b, 90.0, 90.0, 90.0, 1e-2, 1e-2);
}

// Orthorhombic P: all angles 90, unequal edges, P-centering
TEST_CASE("LatticeSearch - orthorhombic P") {
    const double a = 35.0, b = 41.0, c = 57.0;
    CrystalLattice L(a,b,c, 90.0,90.0,90.0);

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(res.centering == 'P');

    auto uc = res.conventional.GetUnitCell();
    check_uc(uc, a, b, c, 90.0, 90.0, 90.0, 1e-6, 1e-4);
}

// Hexagonal P: a=b!=c, alpha=beta=90, gamma=120, P-centering
TEST_CASE("LatticeSearch - hexagonal P") {
    const double a = 30.0, c = 48.0;
    CrystalLattice L(
        Coord(a, 0, 0),
        Coord(-a/2, a*std::sqrt(3)/2, 0),
        Coord(0, 0, c)
    );

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Hexagonal);
    CHECK(res.centering == 'P');

    auto uc = res.conventional.GetUnitCell();
    CHECK(uc.a == Catch::Approx(a).margin(1e-2));
    CHECK(uc.b == Catch::Approx(a).margin(1e-2));
    CHECK(uc.c == Catch::Approx(c).margin(1e-2));
    CHECK(uc.alpha == Catch::Approx(90.0).margin(1e-2));
    CHECK(uc.beta  == Catch::Approx(90.0).margin(1e-2));
    CHECK(uc.gamma == Catch::Approx(120.0).margin(1e-2));
}

TEST_CASE("LatticeSearch - monoclinic C (unique b)") {
    const double a = 50.0, b = 60.0, c = 70.0;
    const double alpha = 90.0, beta = 96.0, gamma = 90.0;
    CrystalLattice conv(a,b,c, alpha,beta,gamma);
    auto L = conv.ToPrimitive('C');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Monoclinic);
    CHECK(res.centering == 'C');

    auto uc = res.conventional.GetUnitCell();
    // Check right angles at alpha,gamma and non-90 beta; lengths comparable
    CHECK(std::fabs(uc.alpha - 90.0) < 1e-3);
    CHECK(std::fabs(uc.gamma - 90.0) < 1e-3);
    CHECK(std::fabs(uc.beta  - beta) < 1e-2);
    // Lengths should match within small tolerance
    CHECK(uc.a == Catch::Approx(a).margin(1e-2));
    CHECK(uc.b == Catch::Approx(b).margin(1e-2));
    CHECK(uc.c == Catch::Approx(c).margin(1e-2));
}

TEST_CASE("LatticeSearch - monoclinic C (unique b) - v2") {
    const double a = 71.0, b = 35.0, c = 90.0;
    const double alpha = 90.0, beta = 96.0, gamma = 90.0;
    CrystalLattice conv(a,b,c, alpha,beta,gamma);
    auto L = conv.ToPrimitive('C');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Monoclinic);
    CHECK(res.centering == 'C');

    auto uc = res.conventional.GetUnitCell();
    // Check right angles at alpha,gamma and non-90 beta; lengths comparable
    CHECK(std::fabs(uc.alpha - 90.0) < 1e-3);
    CHECK(std::fabs(uc.gamma - 90.0) < 1e-3);
    CHECK(std::fabs(uc.beta  - beta) < 1e-2);
    // Lengths should match within small tolerance
    CHECK(uc.a == Catch::Approx(a).margin(1e-2));
    CHECK(uc.b == Catch::Approx(b).margin(1e-2));
    CHECK(uc.c == Catch::Approx(c).margin(1e-2));
}

TEST_CASE("LatticeSearch - monoclinic C (unique a)") {
    const double a = 60.0, b = 50.0, c = 70.0;
    const double alpha = 96.0, beta = 90.0, gamma = 90.0;
    CrystalLattice conv(a,b,c, alpha,beta,gamma);
    auto L = conv.ToPrimitive('C');

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Monoclinic);
    CHECK(res.centering == 'C');

    auto uc = res.conventional.GetUnitCell();
    // Check right angles at alpha,gamma and non-90 beta; lengths comparable
    CHECK(std::fabs(uc.alpha - 90.0) < 1e-3);
    CHECK(std::fabs(uc.gamma - 90.0) < 1e-3);
    CHECK(std::fabs(uc.beta  - alpha) < 1e-2);
    // Lengths should match within small tolerance
    CHECK(uc.a == Catch::Approx(b).margin(1e-2));
    CHECK(uc.b == Catch::Approx(a).margin(1e-2));
    CHECK(uc.c == Catch::Approx(c).margin(1e-2));
}

TEST_CASE("LatticeSearch - monoclinic P (unique b)") {
    const double a = 50.0, b = 60.0, c = 70.0;
    const double alpha = 90.0, beta = 96.0, gamma = 90.0;
    CrystalLattice conv(a,b,c, alpha,beta,gamma);

    auto res = LatticeSearch(conv, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Monoclinic);
    CHECK(res.centering == 'P');

    auto uc = res.conventional.GetUnitCell();
    // Check right angles at alpha,gamma and non-90 beta; lengths comparable
    CHECK(std::fabs(uc.alpha - 90.0) < 1e-3);
    CHECK(std::fabs(uc.gamma - 90.0) < 1e-3);
    CHECK(std::fabs(uc.beta  - beta) < 1e-2);
    // Lengths should match within small tolerance
    CHECK(uc.a == Catch::Approx(a).margin(1e-2));
    CHECK(uc.b == Catch::Approx(b).margin(1e-2));
    CHECK(uc.c == Catch::Approx(c).margin(1e-2));
}

TEST_CASE("LatticeSearch - monoclinic P (unique b) - v2") {
    const double a = 90.0, b = 35.0, c = 71.0;
    const double alpha = 90.0, beta = 96.0, gamma = 90.0;
    CrystalLattice conv(a,b,c, alpha,beta,gamma);

    auto res = LatticeSearch(conv, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Monoclinic);
    CHECK(res.centering == 'P');

    auto uc = res.conventional.GetUnitCell();
    // Check right angles at alpha,gamma and non-90 beta; lengths comparable
    CHECK(std::fabs(uc.alpha - 90.0) < 1e-3);
    CHECK(std::fabs(uc.gamma - 90.0) < 1e-3);
    CHECK(std::fabs(uc.beta  - beta) < 1e-2);
    // Lengths should match within small tolerance
    CHECK(uc.a == Catch::Approx(c).margin(1e-2));
    CHECK(uc.b == Catch::Approx(b).margin(1e-2));
    CHECK(uc.c == Catch::Approx(a).margin(1e-2));
}

TEST_CASE("LatticeSearch - triclinic P") {
    // General triclinic primitive cell
    CrystalLattice L(33.1, 41.7, 52.3, 89.1, 85.0, 76.3);

    auto res = LatticeSearch(L, 1e-6);

    // System should be triclinic, centering P, and conventional equals some standardized primitive
    CHECK(res.system == gemmi::CrystalSystem::Triclinic);
    CHECK(res.centering == 'P');

    // The conventional cell should be metric-equivalent to input. We verify only the system and centering here.
    // Reduced primitive must be non-singular
    auto uc_red = res.primitive_reduced.GetUnitCell();
    CHECK(uc_red.a > 0);
    CHECK(uc_red.b > 0);
    CHECK(uc_red.c > 0);
}

TEST_CASE("LatticeSearch - triclinic P - v2") {
    // General triclinic primitive cell
    CrystalLattice L(33.1, 41.7, 52.3, 100, 92, 115);

    auto res = LatticeSearch(L, 1e-6);

    // System should be triclinic, centering P, and conventional equals some standardized primitive
    CHECK(res.system == gemmi::CrystalSystem::Triclinic);
    CHECK(res.centering == 'P');

    // The conventional cell should be metric-equivalent to input. We verify only the system and centering here.
    // Reduced primitive must be non-singular
    auto uc_red = res.primitive_reduced.GetUnitCell();
    CHECK(uc_red.a > 0);
    CHECK(uc_red.b > 0);
    CHECK(uc_red.c > 0);
}

TEST_CASE("LatticeSearch - trigonal R") {
    const double a = 32.0;
    const double alpha = 80.0;

    // Build rhombohedral in rhombohedral setting (primitive axes a=b=c, alpha=beta=gamma)
    CrystalLattice L(a, a, a, alpha, alpha, alpha);

    auto res = LatticeSearch(L, 1e-6);

    CHECK(res.system == gemmi::CrystalSystem::Trigonal);
    CHECK(res.centering == 'R');

    auto uc_red = res.conventional.GetUnitCell();
    CHECK(uc_red.alpha == Catch::Approx(90).margin(1e-2));
    CHECK(uc_red.beta == Catch::Approx(90).margin(1e-2));
    CHECK(uc_red.gamma == Catch::Approx(120).margin(1e-2));

    auto uc_prim = res.primitive_reduced.GetUnitCell();
    CHECK(uc_prim.alpha == Catch::Approx(alpha).margin(1e-2));
    CHECK(uc_prim.beta == Catch::Approx(alpha).margin(1e-2));
    CHECK(uc_prim.gamma == Catch::Approx(alpha).margin(1e-2));
}

// The class-filtered walk: the same table, restricted to one Bravais class. A tetragonal-P lattice is
// also a C-centred orthorhombic one (a_C = a+b, b_C = -a+b, c_C = c), and asking for that class has to
// return that setting even though the plain search rightly prefers the tetragonal one.
TEST_CASE("LatticeSearchForClass - tetragonal P also has a C-centred orthorhombic setting") {
    const double a = 50.0, c = 120.0;
    const CrystalLattice L(a, a, c, 90, 90, 90);

    const auto plain = LatticeSearch(L, 1e-6);
    CHECK(plain.system == gemmi::CrystalSystem::Tetragonal);
    CHECK(plain.centering == 'P');

    const auto ortho = LatticeSearchForClass(L, gemmi::CrystalSystem::Orthorhombic, 'C', 1e-6);
    REQUIRE(ortho.has_value());
    CHECK(ortho->system == gemmi::CrystalSystem::Orthorhombic);
    CHECK(ortho->centering == 'C');
    const auto uc = ortho->conventional.GetUnitCell();
    // The C cell is the face diagonal on a and b, so twice the volume and a = b = a_tet * sqrt(2).
    CHECK(uc.a == Catch::Approx(a * std::sqrt(2.0)).margin(1e-4));
    CHECK(uc.b == Catch::Approx(a * std::sqrt(2.0)).margin(1e-4));
    CHECK(uc.c == Catch::Approx(c).margin(1e-4));
    CHECK(uc.alpha == Catch::Approx(90).margin(1e-4));
    CHECK(uc.beta == Catch::Approx(90).margin(1e-4));
    CHECK(uc.gamma == Catch::Approx(90).margin(1e-4));
}

TEST_CASE("LatticeSearchForClass - a class the metric cannot carry is refused") {
    // A general triclinic metric has no monoclinic-C setting, and an F-centred cubic lattice has no
    // hexagonal-P one (its hexagonal description is R-centred).
    const CrystalLattice tri(41.0, 47.0, 53.0, 71.0, 83.0, 97.0);
    CHECK_FALSE(LatticeSearchForClass(tri, gemmi::CrystalSystem::Monoclinic, 'C').has_value());

    const double a = 60.0;
    const auto cubic_f = CrystalLattice(a, a, a, 90, 90, 90).ToPrimitive('F');
    CHECK(LatticeSearch(cubic_f, 1e-6).centering == 'F');
    CHECK_FALSE(LatticeSearchForClass(cubic_f, gemmi::CrystalSystem::Hexagonal, 'P').has_value());
    // ... but its rhombohedral setting is there, which is what makes the refusal above a real answer
    // rather than an artefact of the filter.
    const auto rhomb = LatticeSearchForClass(cubic_f, gemmi::CrystalSystem::Trigonal, 'R');
    REQUIRE(rhomb.has_value());
    CHECK(rhomb->centering == 'R');
}

TEST_CASE("LatticeSearchForClass - asking for what the plain search found returns the same setting") {
    const double a = 40.0;
    const auto L = CrystalLattice(a, a, a, 90, 90, 90).ToPrimitive('I');
    const auto plain = LatticeSearch(L, 1e-6);
    const auto filtered = LatticeSearchForClass(L, plain.system, plain.centering, 1e-6);
    REQUIRE(filtered.has_value());
    CHECK(filtered->niggli_class == plain.niggli_class);
    check_uc(filtered->conventional.GetUnitCell(), a, a, a, 90, 90, 90, 1e-4, 1e-4);
}