scene-actor-extraction/src/nodes/identity_matcher_node.hpp

#pragma once
#include "types.hpp"
#include "config.hpp"
#include "inference/similarity.hpp"
#include "gallery/gallery_store.hpp"
#include "gallery/gallery_calibration.hpp"
#include "gallery/track_gallery.hpp"

#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>

// ── IdentityMatcherFunc ───────────────────────────────────────────────────────
// KPN node: compares each embedding against every reference embedding in the
// actor gallery using cosine similarity.
//
// Matching strategy — two modes selected at construction time:
//
//   Calibrated (preferred): gallery calibration fits a sigmoid
//     P(match) = σ(a·similarity + b) from intra/inter-class pairs.
//     A face is accepted if P(match | best_actor) > prob_threshold.
//
//   Fallback (no calibration): dual-criterion accept —
//     (a) best cosine distance < match_threshold, OR
//     (b) ratio test: best_dist/second_best_dist < match_ratio
//         AND best_dist < match_ratio_ceil.
//
// In both modes, per-actor best similarity is determined by scanning
// reference embeddings and taking the closest (best-of-N).
//
// Gallery scan: the full reference set (tens of thousands of 512-dim
// embeddings) is uploaded to the GPU once at construction time and stays
// resident there. Per frame, only the small query matrix (n_faces x 512) is
// uploaded and a single SGEMM computes the full similarity matrix in well under
// a millisecond. The GPU math backend (cuBLAS or rocBLAS) lives behind
// ISimilarityEngine (backends/gemm_backend.cpp) and is selected at compile time.

struct IdentityMatcherFunc {
    static constexpr std::string_view label() { return "identity_matcher"; }

    // Max faces handled per frame without reallocating GPU buffers.
    static constexpr int kMaxFaces = 32;

    IdentityMatcherFunc(const ActorGallery& gallery, const Config& cfg)
        : gallery_(gallery)
        , prob_threshold_(cfg.prob_threshold)
        , log_prior_odds_(std::log(cfg.match_prior / (1.f - cfg.match_prior)))
        , threshold_(cfg.match_threshold)
        , ratio_(cfg.match_ratio)
        , ratio_ceil_(cfg.match_ratio_ceil)
        , track_gallery_(cfg)
    {
        std::cerr << "[identity_matcher] flattening gallery embeddings...\n";
        for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
            for (const auto& emb : gallery_.actors[ai].embeddings) {
                flat_emb_.push_back(emb);
                flat_actor_.push_back(ai);
            }
        }
        n_gallery_ = static_cast<int>(flat_emb_.size());

        std::cerr << "[identity_matcher] starting calibration ("
                  << flat_emb_.size() << " embeddings)...\n";
        bool recomputed = false;
        cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_,
                                        gallery_.calib_a, gallery_.calib_b,
                                        gallery_.calib_valid, gallery_.calib_hash,
                                        cfg.gallery_path + ".calib_cache", recomputed);
        if (recomputed) {
            // Persist the freshly-fitted calibration into the gallery file (always
            // HDF5 — save_gallery rewrites any other extension, see gallery_store.cpp)
            // so the next run against this same, unchanged gallery skips the O(n^2) fit.
            ActorGallery to_save = gallery_;
            to_save.calib_a     = cal_.a;
            to_save.calib_b     = cal_.b;
            to_save.calib_valid = cal_.valid;
            to_save.calib_hash  = hash_gallery_embeddings(flat_emb_, flat_actor_);
            save_gallery(cfg.gallery_path, to_save);
            std::cerr << "[identity_matcher] wrote refreshed calibration back to "
                      << cfg.gallery_path << "\n";
        }

        if (cal_.valid) {
            std::cerr << "[identity_matcher] calibrated Bayesian matching"
                      << "  prior=" << cfg.match_prior
                      << "  P_threshold=" << prob_threshold_
                      << "  effective_sim_boundary="
                      << cal_.boundary_at(prob_threshold_, log_prior_odds_) << "\n";
        } else {
            std::cerr << "[identity_matcher] threshold matching (calibration skipped)"
                      << "  threshold=" << threshold_
                      << "  ratio=" << ratio_ << "  ratio_ceil=" << ratio_ceil_ << "\n";
        }
        std::cerr << "[identity_matcher] gallery: "
                  << gallery_.actors.size() << " actors, "
                  << flat_emb_.size() << " reference embeddings\n";

        std::vector<float> host_gallery(static_cast<size_t>(n_gallery_) * 512);
        for (int i = 0; i < n_gallery_; ++i)
            std::memcpy(host_gallery.data() + static_cast<size_t>(i) * 512,
                        flat_emb_[i].data(), 512 * sizeof(float));

        sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces);
    }

    /// TRACES: AR-023, AR-024 | SR-002
    /// The fitted sigmoid. Exposed because the matcher is where it gets fitted
    /// (and cached back to the gallery), but it is not the matcher's private
    /// property: track association and evidence weighting must threshold in the
    /// *same* probability space, or a "0.5" in one stage and a "0.5" in another
    /// mean different things. See `same_person_probability`.
    const GalleryCalibration& calibration() const { return cal_; }

    // Runtime setter — lets a persistent pipeline be reused across a threshold sweep
    // without rebuilding the (expensive, gallery-resident) matcher. The gallery,
    // calibration and GPU sim-engine stay put; only the accept threshold changes.
    void set_prob_threshold(float t) { prob_threshold_ = t; }

    MatchedSceneFrame operator()(TrackedSceneFrame tf) {
        if (tf.source.eof) {
            track_gallery_.clear_tracks();
            return {std::move(tf.source), {}};
        }

        // A hard cut changes the camera viewpoint. The face_tracker may revive a
        // track_id across the cut (identity continuity), but promotion must never
        // mix embeddings from two viewpoints under one buffer, so we still drop
        // every diversity buffer here — a revived track simply re-accumulates its
        // buffer from post-cut frames. Stale cross-cut embeddings are never promoted.
        if (tf.source.is_cut) track_gallery_.clear_tracks();

        const int n_faces = static_cast<int>(tf.embeddings.size());
        std::vector<IdentifiedActor> actors;
        actors.reserve(n_faces);

        if (n_faces == 0) return {std::move(tf.source), {}};
        if (n_faces > kMaxFaces)
            throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");

        std::vector<float> host_query(static_cast<size_t>(n_faces) * 512);
        for (int fi = 0; fi < n_faces; ++fi) {
            std::memcpy(host_query.data() + static_cast<size_t>(fi) * 512,
                        tf.embeddings[fi].data(), 512 * sizeof(float));
        }

        // S (N_gallery × n_faces) col-major: face fi's gallery sims at sims + fi*n_gallery.
        const float* host_sims = sim_engine_->compute(host_query.data(), n_faces);

        for (int fi = 0; fi < n_faces; ++fi) {
            const float* sims = host_sims + static_cast<size_t>(fi) * n_gallery_;

            std::vector<float> best_sim(gallery_.actors.size(),
                                        -std::numeric_limits<float>::max());
            for (int ei = 0; ei < n_gallery_; ++ei) {
                float sim = sims[ei];
                int   ai  = flat_actor_[ei];
                if (sim > best_sim[ai]) best_sim[ai] = sim;
            }

            // Fold in the per-film annex (CPU-side, tens of embeddings). Promoted
            // pose-varied views compete for best-of-N exactly like baked refs, so
            // a face at a pose the gallery lacked can now win its true actor.
            for (const auto& ae : track_gallery_.annex()) {
                float sim = cosine_similarity(tf.embeddings[fi], ae.emb);
                if (sim > best_sim[ae.actor_idx]) best_sim[ae.actor_idx] = sim;
            }

            int   best_actor   = -1;
            int   second_actor = -1;
            float best_s       = -std::numeric_limits<float>::max();
            float second_s     = -std::numeric_limits<float>::max();
            for (int ai = 0; ai < static_cast<int>(best_sim.size()); ++ai) {
                if (best_sim[ai] > best_s) {
                    second_s     = best_s;
                    second_actor = best_actor;
                    best_s       = best_sim[ai];
                    best_actor   = ai;
                } else if (best_sim[ai] > second_s) {
                    second_s     = best_sim[ai];
                    second_actor = ai;
                }
            }
            (void)second_actor;

            bool accept = false;
            if (best_actor >= 0) {
                if (cal_.valid) {
                    accept = cal_.probability(best_s, log_prior_odds_) > prob_threshold_;
                } else {
                    float best_d   = 1.f - best_s;
                    float second_d = (second_s > -std::numeric_limits<float>::max())
                                     ? 1.f - second_s
                                     : std::numeric_limits<float>::max();
                    bool absolute = best_d < threshold_;
                    bool ratio    = (best_d < ratio_ceil_) &&
                                    (second_d == std::numeric_limits<float>::max() ||
                                     best_d / second_d < ratio_);
                    accept = absolute || ratio;
                }
            }

            IdentifiedActor ia;
            // Map bbox back to original video resolution when dense_scale
            // downscaled the decoded frame (detection/tracking ran downscaled;
            // output bboxes must be in original pixel space).
            ia.bbox     = tf.faces[fi].bbox;
            if (tf.source.bbox_upscale != 1.f) {
                const float s = tf.source.bbox_upscale;
                ia.bbox.x *= s; ia.bbox.y *= s;
                ia.bbox.width *= s; ia.bbox.height *= s;
            }
            ia.crop     = tf.crops[fi];
            ia.track_id = tf.track_ids[fi];

            if (accept) {
                ia.actor_idx   = best_actor;
                ia.name        = gallery_.actors[best_actor].name;
                ia.imdb_id     = gallery_.actors[best_actor].imdb_id;
                ia.tmdb_id     = gallery_.actors[best_actor].tmdb_id;
                ia.jellyfin_id = gallery_.actors[best_actor].jellyfin_id;
                ia.similarity = cal_.valid
                                ? cal_.probability(best_s, log_prior_odds_)
                                : best_s;
            }

            // Feed this face into per-film gallery expansion. best_actor/best_s
            // reflect the actor with the strongest gallery similarity for this
            // face (annex already folded in above); the track's diversity buffer
            // keeps the gallery-far views and promotes them once the track is
            // confirmed. No-op unless --expand-gallery is set.
            track_gallery_.observe(tf.track_ids[fi], tf.embeddings[fi],
                                   best_actor, best_s, accept, tf.crops[fi]);

            actors.push_back(std::move(ia));
        }

        return {std::move(tf.source), std::move(actors)};
    }

private:
    ActorGallery           gallery_;
    GalleryCalibration     cal_;
    float                  prob_threshold_;
    float                  log_prior_odds_;
    float                  threshold_;
    float                  ratio_;
    float                  ratio_ceil_;
    std::vector<Embedding> flat_emb_;
    std::vector<int>       flat_actor_;
    int                    n_gallery_{0};

    std::unique_ptr<ISimilarityEngine> sim_engine_;
    TrackGallery                       track_gallery_;
};