// 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 "../image_analysis/indexing/MultiLatticeSearch.h"

namespace {
    // Rotation-invariant comparison via Gram matrix G = L^T L
    void check_same_cell(const CrystalLattice &a, const CrystalLattice &b, double margin) {
        const Coord av[3] = {a.Vec0(), a.Vec1(), a.Vec2()};
        const Coord bv[3] = {b.Vec0(), b.Vec1(), b.Vec2()};
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                CHECK((av[i] * av[j]) == Catch::Approx(bv[i] * bv[j]).margin(margin));
    }
}

TEST_CASE("MultiLatticeSearch_RecoversRotation") {
    CrystalLattice reference(40, 50, 80, 90, 95, 90);

    // Known rotation: 0.4 rad about a tilted axis
    const Coord axis = Coord(0.3f, 0.9f, 0.1f).Normalize();
    const float angle = 0.4f;
    const RotMatrix R(angle, axis);

    const CrystalLattice rotated(R * reference.Vec0(),
                                 R * reference.Vec1(),
                                 R * reference.Vec2());

    auto result = MultiLatticeSearch({reference, rotated});

    REQUIRE(result.size() == 2);

    // First entry: identity
    CHECK(result[0].rotation_vector.Length() == Catch::Approx(0.0).margin(1e-6));
    check_same_cell(result[0].output_lattice, reference, 1e-3);

    // Second entry: angle and axis recovered
    CHECK(result[1].rotation_vector.Length() == Catch::Approx(angle).margin(1e-4));
    const Coord recovered_axis = result[1].rotation_vector.Normalize();
    CHECK(recovered_axis.x == Catch::Approx(axis.x).margin(1e-3));
    CHECK(recovered_axis.y == Catch::Approx(axis.y).margin(1e-3));
    CHECK(recovered_axis.z == Catch::Approx(axis.z).margin(1e-3));

    // output_lattice is a proper rotation of the reference => same metric
    check_same_cell(result[1].output_lattice, reference, 1e-2);

    // and output equals the rotated input (same orientation)
    check_same_cell(result[1].output_lattice, rotated, 1e-2);
}

TEST_CASE("MultiLatticeSearch_SkipsDifferentCell") {
    CrystalLattice reference(40, 50, 80, 90, 90, 90);
    CrystalLattice other_cell(45, 50, 80, 90, 90, 90); // a differs by 5 A
    

    auto result = MultiLatticeSearch({reference, other_cell});

    // Only the reference survives
    REQUIRE(result.size() == 1);
    CHECK(result[0].rotation_vector.Length() == Catch::Approx(0.0).margin(1e-6));
}

TEST_CASE("MultiLatticeSearch_Empty") {
    auto result = MultiLatticeSearch({});
    CHECK(result.empty());
}

TEST_CASE("MultiLatticeSearch_EP") {
    // Real EP case
    CrystalLattice cell1(Coord(-13.2, -30.0, -29.6),
        Coord(70.1, -12.16, -18.6),
        Coord(7.6, -23.9, 44.8));
    CrystalLattice cell2(Coord(-13.2, -29.9, -29.5),
        Coord(-70.1, 12.15, 18.8),
        Coord(1.8, 45.4, -23.7)
    );

    auto result = MultiLatticeSearch({cell1, cell2},
        0.1, 5);
    REQUIRE(result.size() == 2);
}