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 "track_registry.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 — one mode, always.
//
//   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. Per-actor best similarity is the closest reference
//   embedding (best-of-N).
//
/// TRACES: AR-024 | SR-002
// **There is no raw-cosine fallback.** There used to be: when the fit was
// invalid this node switched to a cosine-distance ceiling plus a ratio test
// (`match_threshold`, `match_ratio`, `match_ratio_ceil`). Three things were
// wrong with it, and the third is the one that mattered.
//
//   1. It violated AR-024 outright, untagged — a bare cosine threshold means
//      something different for every model, gallery and face size.
//   2. It disagreed with the rest of the pipeline about what "calibration
//      failed" means. `same_person_probability` answers that question by
//      falling back to the untuned default sigmoid and saying so loudly, so
//      tracking and evidence weighting stayed in probability space while
//      matching alone left it. One run, two policies.
//   3. Its accepted faces were still fed to `TrackRegistry::observe`, whose
//      contract reads "posterior is a calibrated probability, never a raw
//      cosine (AR-024) ... so the accumulation cannot be fed an uncalibrated
//      number by a careless caller". It could. `max(0, cosine)` went straight
//      into the log-odds accumulation as though it were a probability.
//
// An invalid fit now behaves exactly as everywhere else: the default sigmoid,
// with a warning that says the probabilities are not meaningful. That is a
// worse answer than a fitted calibration and a better one than a number whose
// units nothing else in the pipeline shares.
//
// 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.
//
// TRACES: AR-026 | SR-001
// That resident matrix grows during a film: per-film expansion (AR-019) promotes
// pose-varied views, and they are APPENDED to it rather than scored separately,
// so one multiply covers baked and promoted references alike and best-of-N is a
// single pass over one similarity column. There is no second similarity path in
// this node to fall out of step with the first.

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

        /// TRACES: AR-024 | SR-002
        // Same sentence either way, because it is the same decision rule; only
        // the provenance of (a, b) differs. An unfitted sigmoid still returns
        // plausible-looking probabilities, so the warning has to be the thing
        // that distinguishes them — nothing downstream can.
        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";
        if (!cal_.valid) {
            std::cerr << "[identity_matcher] WARNING: the calibration is NOT fitted "
                         "(a=" << cal_.a << ", b=" << cal_.b << ") — matching runs "
                         "on the untuned default sigmoid, so prob_threshold is not "
                         "comparable to a tuned run's.\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-018, AR-024 | SR-005
        // The expansion store thresholds in the same probability space as
        // association and evidence weighting, so a "0.9" means one thing
        // pipeline-wide rather than three.
        track_gallery_.set_calibration(same_person_probability(cal_));
    }

    /// 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_; }

    /// TRACES: AR-012, AR-025 | SR-002
    /// Where per-frame identity evidence reaches the registry. Optional: with no
    /// registry attached the matcher behaves exactly as before, which keeps the
    /// replay harness and the unit tests working unchanged.
    void set_registry(std::shared_ptr<TrackRegistry> r) {
        registry_ = std::move(r);
        // This node is the evidence source, so the registry must not close a
        // track until this node's watermark has passed it (AR-013).
        if (registry_) registry_->expect_evidence();
    }

    // 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), {}};
        }

        /// TRACES: AR-012, AR-013 | SR-002
        // Publish the evidence watermark BEFORE voting on this frame: every
        // observation strictly before it has now been folded in, so the registry
        // may reap against it. Unconditional -- a frame with no faces still
        // advances the watermark, or a long faceless stretch would stall reaping
        // and hold every dormant track open to the end of the film.
        //
        // This is what makes presence independent of node speed. The registry
        // used to reap on the TRACKER's clock, and backpressure (working as
        // AR-004 intends) means the tracker can be a whole channel's depth ahead
        // of this node -- so tracks were closed before their votes arrived, the
        // votes were dropped, and the run silently under-reported. Measured on
        // the SuperHero fixture before this change: channel depth 32 gave 5
        // actors, depth 10322 gave 0, from identical input.
        if (registry_) registry_->advance_evidence(tf.source.timestamp_sec);

        // 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.
        /// TRACES: AR-019 | SR-005
        // Promotion may only borrow same-identity evidence from a span where
        // identity is certain, so ALL THREE discontinuity signals clear the
        // buffers, not just the histogram cut:
        //   is_cut            — camera-angle change
        //   is_scene_boundary — different scene (AR-010; previously never set,
        //                       so this half of the gate was dead)
        // The third, an identity contradiction (AR-015), is enforced by the
        // registry: a track whose belief swapped is closed outright, so it can
        // no longer promote anything.
        if (tf.source.is_cut || tf.source.is_scene_boundary)
            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), {}};

        /// TRACES: AR-003, AR-004 | SR-002
        // kMaxFaces sizes the similarity engine's preallocated buffer, so it
        // bounds MEMORY, not how many faces a frame may contain. It used to
        // throw above the bound, which made it a hard cap on crowd scenes by
        // accident; now the frame is scored in batches of that size.
        //
        // Faces per frame are unbounded (AR-003) because X-Ray credits scene
        // membership to background cast too, and a fixed cap discards exactly
        // those — the smallest faces are dropped first. Cost is contained by
        // backpressure (AR-004), which slows the producer, rather than by
        // silently throwing work away.
        std::vector<float> host_query(static_cast<size_t>(kMaxFaces) * 512);

        for (int base = 0; base < n_faces; base += kMaxFaces) {
        const int chunk = std::min(kMaxFaces, n_faces - base);
        for (int k = 0; k < chunk; ++k) {
            std::memcpy(host_query.data() + static_cast<size_t>(k) * 512,
                        tf.embeddings[base + k].data(), 512 * sizeof(float));
        }

        /// TRACES: AR-026 | SR-001
        // One GEMM now covers baked references AND the per-film annex: promoted
        // rows were appended to the engine's resident matrix, so they are just
        // more gallery rows with an entry in flat_actor_. The annex used to be
        // folded in afterwards by a host-side cosine loop, justified by "tens of
        // embeddings" — an assumption AR-018/AR-019 retired, since every owned
        // track promotes and the annex grows with cast size and film length.
        //
        // n_gallery() is read per frame, not cached: it grows as promotions land.
        const int    n_gal     = sim_engine_->n_gallery();
        const float* host_sims = sim_engine_->compute(host_query.data(), chunk);

        for (int ci = 0; ci < chunk; ++ci) {
            const int fi = base + ci;
            const float* sims = host_sims + static_cast<size_t>(ci) * n_gal;

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

            // Only the best matters now. The runner-up was tracked solely for
            // the retired ratio test, which asked whether the best cosine stood
            // out from the second — a question the calibrated posterior does
            // not need, since it already says how likely the best match is to
            // be right rather than how much it beat its neighbour by.
            int   best_actor = -1;
            float best_s     = -std::numeric_limits<float>::max();
            for (int ai = 0; ai < static_cast<int>(best_sim.size()); ++ai) {
                if (best_sim[ai] > best_s) {
                    best_s     = best_sim[ai];
                    best_actor = ai;
                }
            }

            /// TRACES: AR-024 | SR-002
            // One rule, whatever the fit's provenance. The cosine reaches a
            // comparison only through cal_.probability().
            const float best_p = best_actor >= 0
                               ? cal_.probability(best_s, log_prior_odds_)
                               : 0.f;
            const bool accept = best_actor >= 0 && best_p > prob_threshold_;

            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 = best_p;
            }

            // 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.
            // TRACES: AR-012, AR-025 | SR-002
            // Every scored face is evidence, not only the accepted ones: a run of
            // near-misses for one actor is itself informative, and discarding it
            // would make ownership depend on a per-frame threshold the redesign
            // exists to stop relying on. The registry discounts for correlation
            // and decides ownership from the accumulated posterior (AR-025).
            if (registry_ && best_actor >= 0 && tf.track_ids[fi] >= 0)
                registry_->observe(tf.track_ids[fi], best_actor, best_p,
                                   tf.embeddings[fi]);

            // TRACES: AR-019 | SR-005
            // Ownership is the registry's, computed once. TrackGallery used to
            // tally its own plurality vote over accepted frames, which meant two
            // different answers to "who is this track" could coexist — and the
            // expansion one ignored the Bayesian accumulation entirely.
            if (registry_ && tf.track_ids[fi] >= 0) {
                if (auto owner = registry_->owner(tf.track_ids[fi]))
                    track_gallery_.set_owner(tf.track_ids[fi], *owner);
            }

            track_gallery_.observe(tf.track_ids[fi], tf.embeddings[fi],
                                   best_actor, best_s, accept, tf.crops[fi]);

            actors.push_back(std::move(ia));
        }
        }   // chunk loop

        absorb_promotions();

        /// TRACES: AR-019 | SR-005
        // Drop buffers for tracks the registry has reaped. Without this a track
        // that simply went off screen kept its diversity buffer until the next
        // cut, so the store grew with the film rather than with what is on
        // screen — and a buffer that outlives its track is evidence about a
        // person nobody is looking at any more.
        if (registry_)
            track_gallery_.prune_dead(
                [this](int id) { return registry_->is_live(id); });

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

private:
    /// TRACES: AR-026 | SR-001
    /// Move rows promoted during this frame into the resident gallery matrix,
    /// extending the actor mapping in lockstep so row i keeps naming the actor
    /// at flat_actor_[i]. Runs once per frame, after every face has been scored:
    /// appending mid-frame would invalidate the similarity pointer the chunk
    /// loop is still reading, and it is also the semantics the expansion store
    /// documents — a promotion helps SUBSEQUENT frames, never the one that
    /// produced it, so identification cannot depend on face order within a frame.
    void absorb_promotions() {
        if (!track_gallery_.enabled()) return;

        pending_emb_.clear();
        pending_actor_.clear();
        const int n = track_gallery_.drain_promotions(pending_emb_, pending_actor_);
        if (n == 0) return;

        sim_engine_->append_rows(pending_emb_.data(), n);
        flat_actor_.insert(flat_actor_.end(),
                           pending_actor_.begin(), pending_actor_.end());
        n_gallery_ = sim_engine_->n_gallery();
    }

    ActorGallery           gallery_;
    GalleryCalibration     cal_;
    float                  prob_threshold_;
    float                  log_prior_odds_;
    /// flat_emb_ is the BAKED reference set only — it is the calibration fit's
    /// input (AR-023) and is not touched again after construction. flat_actor_,
    /// by contrast, is the actor mapping parallel to the *engine's* rows, so it
    /// grows with every promotion absorbed (AR-026) and is the longer of the two.
    std::vector<Embedding> flat_emb_;
    std::vector<int>       flat_actor_;
    int                    n_gallery_{0};

    // Reused across frames so absorbing a promotion allocates nothing.
    std::vector<float>     pending_emb_;
    std::vector<int>       pending_actor_;

    std::unique_ptr<ISimilarityEngine> sim_engine_;
    TrackGallery                       track_gallery_;
    std::shared_ptr<TrackRegistry>     registry_;
};