// TRACES: UT-004 | AR-026 | SR-001
//
// 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.
//
// This is the CI half of AR-026: equivalence between the GEMM path and
// hand-computed dot products on small input. Throughput at scale (AR-027) is T4
// and cannot run here.
#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));
}

// TRACES: UT-004 | AR-026 | SR-001
// The annex half of AR-026: promoted rows are appended to the resident matrix
// and scored by the same GEMM as the baked references. What used to be a
// host-side cosine loop over TrackGallery::annex() is now these extra columns,
// so the equivalence that matters is "an appended row scores exactly what the
// reference dot product says", and "appending changes nothing about the rows
// already there".
TEST_CASE("appended rows are scored by the same GEMM as the baked gallery",
          "[similarity][AR-026]") {
    std::vector<float> gallery;
    for (int slot : {0, 1}) {
        auto e = one_hot(slot);
        gallery.insert(gallery.end(), e.begin(), e.end());
    }
    auto engine = make_similarity_engine(gallery.data(), /*n_gallery=*/2, /*max_faces=*/2);
    REQUIRE(engine->n_gallery() == 2);

    // One query face at 45° between slots 1 and 2. Slot 2 is not in the baked
    // gallery yet, so the face is currently "unrecognised at that pose".
    const float s = std::sqrt(0.5f);
    std::vector<float> query(static_cast<size_t>(2) * 512, 0.0f);
    query[1] = s;
    query[2] = s;

    const float* before = engine->compute(query.data(), 1);
    CHECK_THAT(before[0], WithinAbs(0.0f, 1e-6f));   // vs slot 0
    CHECK_THAT(before[1], WithinAbs(s, 1e-6f));      // vs slot 1

    // Promote the missing view — the annex row an owned track would contribute.
    auto promoted = one_hot(2);
    engine->append_rows(promoted.data(), 1);
    REQUIRE(engine->n_gallery() == 3);

    const float* after = engine->compute(query.data(), 1);
    CHECK_THAT(after[0], WithinAbs(0.0f, 1e-6f));    // baked rows unchanged
    CHECK_THAT(after[1], WithinAbs(s, 1e-6f));
    CHECK_THAT(after[2], WithinAbs(s, 1e-6f));       // appended row, same multiply
}

TEST_CASE("appending many rows keeps every similarity exact", "[similarity][AR-026]") {
    // Start from a one-row gallery and append past the initial capacity several
    // times over — the growth path has to preserve what is already resident, and
    // a film promotes far more rows than the gallery starts with.
    auto seed = one_hot(0);
    auto engine = make_similarity_engine(seed.data(), /*n_gallery=*/1, /*max_faces=*/1);

    constexpr int kAppended = 40;
    for (int i = 1; i <= kAppended; ++i) {
        auto e = one_hot(i);
        engine->append_rows(e.data(), 1);
    }
    REQUIRE(engine->n_gallery() == kAppended + 1);

    // Query one-hot slot k: similarity is 1 against row k and 0 against all others.
    for (int k : {0, 1, 17, kAppended}) {
        auto q = one_hot(k);
        const float* S = engine->compute(q.data(), 1);
        for (int g = 0; g <= kAppended; ++g)
            CHECK_THAT(S[g], WithinAbs(g == k ? 1.0f : 0.0f, 1e-6f));
    }
}

TEST_CASE("appending a block of rows matches appending them one at a time",
          "[similarity][AR-026]") {
    // A promotion hands over a whole diversity buffer at once; that must be
    // indistinguishable from the same rows arriving singly.
    auto seed = one_hot(0);

    std::vector<float> block;
    for (int slot : {1, 2, 3}) {
        auto e = one_hot(slot);
        block.insert(block.end(), e.begin(), e.end());
    }

    auto bulk = make_similarity_engine(seed.data(), 1, 1);
    bulk->append_rows(block.data(), 3);

    auto singly = make_similarity_engine(seed.data(), 1, 1);
    for (int i = 0; i < 3; ++i) singly->append_rows(block.data() + i * 512, 1);

    REQUIRE(bulk->n_gallery() == singly->n_gallery());

    const float t = std::sqrt(1.0f / 3.0f);
    std::array<float, 512> q{};
    q[1] = t; q[2] = t; q[3] = t;

    const float* a = bulk->compute(q.data(), 1);
    std::vector<float> a_copy(a, a + bulk->n_gallery());
    const float* b = singly->compute(q.data(), 1);

    for (int g = 0; g < bulk->n_gallery(); ++g)
        CHECK_THAT(a_copy[g], WithinAbs(b[g], 1e-6f));
}

TEST_CASE("appending zero rows is a no-op", "[similarity][AR-026]") {
    auto e = one_hot(0);
    auto engine = make_similarity_engine(e.data(), 1, 1);
    CHECK_NOTHROW(engine->append_rows(nullptr, 0));
    CHECK(engine->n_gallery() == 1);
}

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));
}