// Unit tests for the CPU reference similarity engine (backends/gemm_backend.cpp,
// SAE_GEMM_CPU) and the l2_normalise helper. All pure, GPU-free, model-free.
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include "face_utils.hpp"
#include "inference/similarity.hpp"

#include <array>
#include <cmath>
#include <vector>

using Catch::Matchers::WithinAbs;

namespace {

// A 512-d embedding that is 1.0 in one slot and 0 elsewhere (already unit-norm).
std::array<float, 512> one_hot(int slot) {
    std::array<float, 512> e{};
    e[slot] = 1.0f;
    return e;
}

}  // namespace

TEST_CASE("l2_normalise produces a unit vector", "[similarity]") {
    std::array<float, 512> raw{};
    raw[0] = 3.0f;
    raw[1] = 4.0f;  // norm 5

    Embedding n = l2_normalise(raw.data());
    CHECK_THAT(n[0], WithinAbs(0.6f, 1e-6f));
    CHECK_THAT(n[1], WithinAbs(0.8f, 1e-6f));

    float norm = 0.f;
    for (float v : n) norm += v * v;
    CHECK_THAT(std::sqrt(norm), WithinAbs(1.0f, 1e-6f));
}

TEST_CASE("l2_normalise guards against a zero vector", "[similarity]") {
    std::array<float, 512> zero{};
    Embedding n = l2_normalise(zero.data());
    for (float v : n) CHECK(v == 0.0f);  // 0 / 1e-6 == 0, no NaN
}

TEST_CASE("CPU similarity engine matches hand-computed dot products", "[similarity]") {
    // Gallery of three orthonormal one-hot embeddings.
    std::vector<float> gallery;
    for (int slot : {0, 1, 2}) {
        auto e = one_hot(slot);
        gallery.insert(gallery.end(), e.begin(), e.end());
    }
    const int n_gallery = 3;
    const int max_faces = 2;

    auto engine = make_similarity_engine(gallery.data(), n_gallery, max_faces);
    REQUIRE(engine->max_faces() == max_faces);

    // Two query faces: face0 == gallery row 1, face1 is 45° between rows 0 and 2.
    std::vector<float> query(static_cast<size_t>(max_faces) * 512, 0.0f);
    query[1] = 1.0f;                       // face0: one-hot slot 1
    const float s = std::sqrt(0.5f);
    query[512 + 0] = s;                    // face1: (1/√2, 0, 1/√2, …)
    query[512 + 2] = s;

    const float* S = engine->compute(query.data(), 2);

    // Column-major: S[g + f*n_gallery].
    // face0 vs gallery {0,1,2} → {0, 1, 0}
    CHECK_THAT(S[0 + 0 * n_gallery], WithinAbs(0.0f, 1e-6f));
    CHECK_THAT(S[1 + 0 * n_gallery], WithinAbs(1.0f, 1e-6f));
    CHECK_THAT(S[2 + 0 * n_gallery], WithinAbs(0.0f, 1e-6f));
    // face1 vs gallery {0,1,2} → {1/√2, 0, 1/√2}
    CHECK_THAT(S[0 + 1 * n_gallery], WithinAbs(s, 1e-6f));
    CHECK_THAT(S[1 + 1 * n_gallery], WithinAbs(0.0f, 1e-6f));
    CHECK_THAT(S[2 + 1 * n_gallery], WithinAbs(s, 1e-6f));
}

TEST_CASE("CPU similarity engine rejects too many faces", "[similarity]") {
    auto e = one_hot(0);
    auto engine = make_similarity_engine(e.data(), /*n_gallery=*/1, /*max_faces=*/1);
    std::array<float, 512 * 2> q{};
    CHECK_THROWS(engine->compute(q.data(), 2));
}

TEST_CASE("CPU similarity engine handles zero query faces", "[similarity]") {
    auto e = one_hot(0);
    auto engine = make_similarity_engine(e.data(), 1, 4);
    // n_faces == 0 must not read the (null) query pointer.
    CHECK_NOTHROW(engine->compute(nullptr, 0));
}