#pragma once
/// TRACES: AR-012, AR-013, AR-014, AR-015, AR-016, AR-017 | SR-002
///
/// TrackRegistry — the single owner of track state and of presence.
///
/// Presence follows **track extent**, not per-frame recognition (AR-012): a
/// window is `[first_seen, last_seen]` of a track an actor owns, so it starts
/// when the actor appeared rather than when the recogniser first succeeded.
///
/// `last_seen` carries the entire liveness state (AR-013):
///
///   unset  → on screen now
///   set    → went off screen at that timestamp, still revivable
///   reaped → emitted to the aggregator and erased
///
/// There is no missing-frame counter and no expired flag; the optional *is* the
/// state machine, and it subsumes what was previously a two-pool split in the
/// tracker (active vs. parked-across-a-cut).
///
/// **Interior gaps are claimed, the trailing cool-down is not.** A face lost at
/// t1 and re-associated at t2 within the timeout never closed its track, so the
/// actor is present across [t1, t2] — correct, since someone briefly occluded or
/// off-camera has not left the scene. But a track that dies ends its window at
/// `last_seen`, never at the moment of death. That asymmetry is what removes the
/// over-claim the retired `extinction_sec` keep-alive produced.
///
/// The registry is created in `main` and shared by `shared_ptr`; it is *not* a
/// KPN node. Ownership is not a stage in the stream — it is state several stages
/// read and write, whose final answer is only known when a track dies.

#include "types.hpp"
#include "evidence_discount.hpp"

#include <cmath>
#include <cstdint>
#include <functional>
#include <map>
#include <mutex>
#include <optional>
#include <string>
#include <vector>

// ── DeadTrack ────────────────────────────────────────────────────────────────
// A finished presence claim, emitted exactly once when a track is reaped or
// flushed. Immutable by construction: it carries everything needed to justify
// itself (AR-017), with no back-reference into registry state.
/// TRACES: AR-017 | IR-002 | SR-002, SR-003
/// How an actor came to be attached to a track.
///
/// AR-017 requires every presence claim to carry its identification route, and
/// IR-002 publishes it per window. Until now the sink wrote the string "live"
/// unconditionally, so the field existed but could not distinguish anything --
/// and AR-017's own verification asks for "deferred and pooled routes
/// distinguishable".
///
/// Only `live` occurs today. `deferred` is what AR-020's pass will set when it
/// resolves a track that failed during streaming and was identified against the
/// final expanded gallery; the value exists now so that pass has somewhere to
/// write rather than a serialisation change to make.
enum class Route {
    live,       ///< identified while streaming, from accumulated per-frame evidence
    deferred,   ///< resolved after EOF against the expanded gallery (AR-020)
};

inline const char* route_name(Route r) {
    switch (r) {
        case Route::deferred: return "deferred";
        case Route::live:     break;
    }
    return "live";
}

struct DeadTrack {
    int    track_id{-1};
    double first_seen{0.0};
    double last_seen{0.0};   ///< always the last sighting, never the death time
    int    actor_idx{-1};    ///< -1 when the track was never owned
    float  belief{0.0f};     ///< accumulated posterior for actor_idx
    Route  route{Route::live}; ///< how the actor was attached (AR-017)
    int    observations{0};    ///< evidence updates that landed on this track
    float  effective_obs{0.f}; ///< sum of discounted weights — the evidence that counted
};

// ── Track ────────────────────────────────────────────────────────────────────
struct Track {
    int                   id{-1};
    double                first_seen{0.0};
    std::optional<double> last_seen;              ///< unset ⇒ on screen
    std::optional<int>    actor;                  ///< set once a posterior crosses
    /// actor_idx → accumulated log(1 − P). Lazy-OR (noisy-OR) accumulation:
    /// each frame is new evidence that this track is that actor, and the
    /// combined belief is the probability that *at least one* sighting was
    /// right. Stored as log(1−P) because that makes the update additive and
    /// keeps precision where it matters — as P approaches 1, (1−P) is the
    /// quantity with the significant digits.
    std::map<int, float>  belief;
    Embedding             mean{};                 ///< running directional mean
    std::vector<Embedding> views;                 ///< distinct looks, for AR-025 discounting
    float                 discounted_weight{0.f}; ///< sum of applied weights
    int                   n_obs{0};                ///< every scored face on this track
    /// Observations that were actually evidence, and so spent the correlation
    /// budget. Indexing the effective-sample correction by this rather than by
    /// n_obs is what stops non-matches exhausting it — see
    /// Config::evidence_floor_p.
    int                   n_evidence{0};

    bool on_screen() const { return !last_seen.has_value(); }
};

// ── TrackRegistry ────────────────────────────────────────────────────────────
class TrackRegistry {
public:
    using DeadTrackFn = std::function<void(const DeadTrack&)>;

    struct Config {
        /// How long a lost track stays available for re-association.
        ///
        /// Named to match Config::track_extinction_sec, which feeds it, and
        /// deliberately NOT `extinction_sec`: that name belonged to the
        /// withdrawn actor keep-alive, and SPEC.md's removal list ends "grep
        /// for both names and expect no survivors". A survivor here would be
        /// the one false positive in that grep, on a field that means
        /// something else entirely -- this one bounds re-association and never
        /// extends a presence claim.
        double track_extinction_sec{5.0};
        float  ownership_logodds{2.0f};   ///< belief needed to own a track (~0.88 posterior)

        /// TRACES: AR-025 | SR-002
        /// Posterior below which an observation is not evidence *for* an actor,
        /// and so does not spend that actor's correlation budget.
        ///
        /// The budget is an effective-sample correction: with observations
        /// correlated at rho, the weight of the n-th is
        /// `n_eff(n+1) - n_eff(n) = 2/((n+1)(n+2))` at rho=0.5, so it decays
        /// quadratically and the total converges to 1/rho = 2. That is the
        /// intended behaviour — a long static shot must not out-argue varied
        /// evidence purely by lasting longer.
        ///
        /// What was not intended is *who spends it*. Every scored face was
        /// folded in, so an observation at p=0.02 — which contributes
        /// log(0.98) = -0.02 of belief, nothing — consumed the same increment
        /// as one at p=0.95. On SuperHero-2 track 3 that exhausted the budget
        /// on the frames that recognised nobody: 103 observations, effective
        /// weight 2.026, belief 0.455 against a 0.881 threshold, with the 51
        /// frames that did identify the actor arriving when each was worth
        /// 0.0002. The identification was lost.
        ///
        /// It also made the answer depend on frame rate, which is the defect
        /// AR-013 already had to fix once: deliver more frames, dilute the
        /// budget with more non-matches, and a track that was owned stops
        /// being owned. Measured — the same clip identified the actor before a
        /// KPN throughput fix and not after, from identical input.
        ///
        /// 0.5 is the point where the posterior stops favouring the hypothesis
        /// at all, not a tuned threshold. Near-misses still count, which is the
        /// design: an observation at 0.6 is evidence and is folded in. Below
        /// 0.5 the observation argues *against*, which noisy-OR cannot
        /// represent, so nothing is lost by declining to spend a budget on it.
        float  evidence_floor_p{0.5f};
    };

    /// The discounter is a constructor argument rather than an option: there is
    /// no correct way to accumulate per-frame evidence without it.
    TrackRegistry(Config cfg, EvidenceDiscounter discounter)
        : cfg_(cfg), discounter_(std::move(discounter)) {}

    void on_track_dead(DeadTrackFn fn) { on_dead_ = std::move(fn); }

    // ── Frame scope ──────────────────────────────────────────────────────────
    // The tracker mutates registry state across a whole association pass, so
    // that pass must be atomic as a unit — per-call locking would let another
    // thread observe a half-updated frame. FrameScope holds the lock for its
    // lifetime and exposes the mutating operations without re-locking.
    class FrameScope {
    public:
        FrameScope(TrackRegistry& reg, double now)
            : reg_(reg), lock_(reg.mu_) { reg_.tick_locked(now); }

        /// All ASSOCIABLE tracks — **one pool**. `last_seen` tells the caller
        /// whether IoU is meaningful; a dormant track is matched on embedding
        /// alone. There is no separate revival path (AR-008).
        ///
        /// TRACES: AR-008, AR-013 | SR-002
        /// Association and reaping share ONE clock — the evidence watermark when a
        /// matcher is attached, the tracker clock otherwise (they coincide when
        /// there is only one). `candidates()` and `reap_locked()` apply the SAME
        /// `track_extinction_sec` horizon against that clock, so the offered pool
        /// and the live pool are the same set:
        ///
        ///   offered      ⟺ (clock - last_seen) ≤ track_extinction_sec
        ///   reaped/erased ⟺ (clock - last_seen) >  track_extinction_sec
        ///
        /// This closes two symmetric failures. (1) Offering on the tracker's clock
        /// (ahead of the watermark) let a face associate onto a track the registry
        /// had ALREADY reaped on the watermark; the vote then landed on a dead id
        /// and was dropped (record_vote → dropped_votes_). Rare live (small lag),
        /// but replay runs the tracker far ahead of the matcher and lost ~0.3% of
        /// votes. (2) Historically, offering on a LOOSER horizon than the reap left
        /// retired tracks in the pool while the matcher lagged, so a new face
        /// re-associated onto a long-dead track and two people merged into one
        /// window (measured: 5 actors/16 windows at depth 32 vs 3/5 at depth 10322).
        /// A single clock and a single threshold make both impossible: nothing is
        /// offered past its reap horizon, nothing is reaped while still offerable.
        std::vector<Track*> candidates() {
            std::vector<Track*> out;
            out.reserve(reg_.tracks_.size());
            // Filter association on the SAME clock reaping uses (the evidence
            // watermark when a matcher is attached, else the tracker clock). The
            // two used to differ deliberately — the tracker offered on now_ while
            // the registry reaped on evidence_through_ — but that let the tracker
            // associate a face onto a track the registry had already reaped on the
            // watermark, whose vote then landed on a dead id and was dropped
            // (record_vote → dropped_votes_). In the live pipeline the lag is tiny
            // so it rarely bit; in replay the Python source runs the tracker far
            // ahead of the matcher and ~0.3% of votes were lost. One clock for both
            // "may this associate?" and "is this reaped?" closes the race: a track
            // past the horizon is neither offered nor reaped-out-from-under a vote.
            const double clock =
                reg_.awaits_evidence_ ? reg_.evidence_through_ : reg_.now_;
            for (auto& [id, t] : reg_.tracks_) {
                // On-screen tracks are always candidates (actively tracked this
                // frame). A dormant (off-screen) track is only worth keeping alive
                // for re-association if it was actually IDENTIFIED: an unowned
                // dormant track has no actor to re-attach to, so holding it in the
                // pool only bloats the matcher's per-frame comparison set (every
                // candidate is a GEMM row) and invites a new face re-associating
                // onto an anonymous stub. Gating dormant tracks on t.actor keeps
                // the pool bounded regardless of how large track_extinction_sec is
                // — which is what makes a long re-association window affordable.
                if (t.last_seen) {   // dormant
                    if (!t.actor.has_value())
                        continue;    // never identified: not worth re-associating
                    if ((clock - *t.last_seen) > reg_.cfg_.track_extinction_sec)
                        continue;    // past the re-association horizon
                }
                out.push_back(&t);
            }
            return out;
        }

        int  create(double t, const Embedding& e)          { return reg_.create_locked(t, e); }
        void mark_seen(int id, double t, const Embedding& e){ reg_.mark_seen_locked(id, t, e); }
        void mark_lost(int id, double last_on_screen)      { reg_.mark_lost_locked(id, last_on_screen); }

    private:
        TrackRegistry&               reg_;
        std::unique_lock<std::mutex> lock_;
    };

    FrameScope begin_frame(double now) { return FrameScope(*this, now); }

    /// Advance the clock and reap. Called every sampled frame **whether or not
    /// it had detections** — without it a track only dies when some other face
    /// happens to appear, and a film ending mid-track never closes.
    void tick(double now) { std::lock_guard g(mu_); tick_locked(now); }

    /// TRACES: AR-012, AR-013, AR-025 | SR-002
    /// The evidence watermark: every observation up to `t` has been folded in.
    ///
    /// Reaping is driven by THIS, not by the tracker's clock, and the difference
    /// is what stops a correct answer from depending on how fast two nodes run.
    ///
    /// The tracker and the matcher are separate KPN nodes with a channel between
    /// them, and the matcher is much the slower of the pair. Backpressure —
    /// working exactly as AR-004 intends — turns that channel's depth into lag,
    /// so the tracker's timestamp can be far ahead of the last frame anybody has
    /// actually voted on. Reaping on the tracker's clock therefore closed tracks
    /// before their evidence arrived: the votes landed on ids that no longer
    /// existed, were counted as dropped, and the track was emitted unowned or
    /// not at all. Deeper channel, fewer identifications, from identical input.
    ///
    /// The fix is not to bound the channel against `track_extinction_sec`. That
    /// makes an algorithm constant police a throughput knob, and leaves the
    /// answer a function of scheduling. It is to reap on the watermark, which is
    /// the same device `SceneBoundaries::scored_through()` uses for the AR-010
    /// join: a consumer past that point is asking about frames nobody has looked
    /// at yet, and the honest response is to wait rather than to guess.
    ///
    /// Monotonic, and only ever *delays* a reap, so no window can be extended by
    /// it — AR-013's "a window ends at the last sighting, never after" is a
    /// property of `emit_locked`, which takes `last_seen` and never `now`.
    void advance_evidence(double t) {
        std::lock_guard g(mu_);
        if (t > evidence_through_) evidence_through_ = t;
        reap_locked();
    }

    /// TRACES: AR-013, AR-025 | SR-002
    /// Declare that some stage will publish an evidence watermark, so reaping
    /// must wait for it.
    ///
    /// Explicit rather than inferred from "has anyone voted yet". Inferring it
    /// re-opens the bug exactly at startup: before the matcher's first frame no
    /// vote has been seen, so the registry would fall back to the tracker's
    /// clock during precisely the window in which the tracker is furthest
    /// ahead. `IdentityMatcherFunc::set_registry` calls this, so any pipeline
    /// with a matcher waits, and a test that drives the tracker alone keeps the
    /// simple behaviour instead of hanging on a watermark nobody will publish.
    void expect_evidence() { std::lock_guard g(mu_); awaits_evidence_ = true; }

    // ── Evidence ─────────────────────────────────────────────────────────────
    /// Fold one observation into a track's belief (AR-025).
    ///
    /// `posterior` is a **calibrated probability**, never a raw cosine
    /// (AR-024) — the registry converts it to log-odds itself, so the
    /// accumulation cannot be fed an uncalibrated number by a careless caller.
    ///
    /// Correlation discounting is applied **here**, not by the caller.
    /// Consecutive frames of one track are near-identical, and accumulating
    /// them as independent evidence drives the posterior to certainty on what
    /// is effectively a single measurement. Leaving that to callers would mean
    /// a forgotten or doubly-applied discount produces confident wrong answers
    /// silently; the registry is the one place all evidence converges, so it is
    /// the one place the correction belongs.
    ///
    /// A vote for a track that has already been reaped is dropped and counted:
    /// a nonzero `dropped_votes()` means the timeout is shorter than the
    /// matcher's lag, which is a real misconfiguration and must not be silent.
    void observe(int track_id, int actor_idx, float posterior, const Embedding& e) {
        std::lock_guard g(mu_);
        auto it = tracks_.find(track_id);
        if (it == tracks_.end()) { ++dropped_votes_; return; }

        Track& t = it->second;
        ++t.n_obs;   // every scored face is seen, whether or not it is evidence

        // Not evidence *for* this actor: contributes ~nothing to the belief and
        // must not spend the correlation budget. See Config::evidence_floor_p.
        if (posterior < cfg_.evidence_floor_p) return;

        const float w = discounter_.weight(t.views, t.n_evidence, e);

        // Weighted lazy-OR: P_new = 1 − (1 − P_old)·(1 − p)^w, which in log
        // space is a plain sum. w is the discounted evidence (AR-025), so a
        // repeated view still advances the belief but by a fraction of what a
        // genuinely new look would.
        const float p = std::min(1.f - 1e-6f, std::max(1e-6f, posterior));
        t.belief[actor_idx] += w * std::log(1.f - p);
        t.discounted_weight += w;
        ++t.n_evidence;

        const int   best = argmax_belief(t);
        const float best_p = 1.f - std::exp(t.belief[best]);
        if (best_p < own_threshold()) return;

        if (!t.actor.has_value()) {
            claim_locked(t, best);
            return;
        }
        if (*t.actor != best) {
            // AR-014 — belief swapped A→B. Not a correction: a track_id almost
            // certainly carried across a viewpoint change onto a different
            // person. Two non-twins both clearing the threshold on one face is
            // not realistic; a track spanning two people is. Continuing would
            // emit one window blending both, so close here and start afresh.
            split_locked(t, best);
        }
    }

    /// Snapshot read: tally and verdict under one lock. Reading them separately
    /// would let a track be both unowned and owned within a single promotion
    /// decision, since the matcher may be voting concurrently.
    std::optional<int> owner(int track_id) const {
        std::lock_guard g(mu_);
        auto it = tracks_.find(track_id);
        return it == tracks_.end() ? std::nullopt : it->second.actor;
    }

    // ── Termination ──────────────────────────────────────────────────────────
    /// Emit every still-live track and empty the registry (AR-016). A film ends
    /// with faces on screen and those tracks have not timed out, so without this
    /// the closing scene's cast is silently never emitted — a loss that presents
    /// as a recognition miss rather than a bookkeeping bug.
    ///
    /// Idempotent: calling it twice emits nothing the second time.
    void flush(double final_ts) {
        std::lock_guard g(mu_);
        for (auto& [id, t] : tracks_) emit_locked(t, t.last_seen.value_or(final_ts));
        tracks_.clear();
    }

    // ── Diagnostics ──────────────────────────────────────────────────────────
    // These measure how often tracking is silently wrong, which nothing in the
    // pipeline currently reveals.
    /// Whether the registry still holds this track. The authority on which
    /// tracks exist, so annotating structures elsewhere (spatial boxes in the
    /// tracker, diversity buffers in the expansion store) can prune against it
    /// rather than keeping a second opinion.
    bool is_live(int track_id) const {
        std::lock_guard g(mu_);
        return tracks_.count(track_id) != 0;
    }

    int dropped_votes()   const { std::lock_guard g(mu_); return dropped_votes_; }
    int belief_swaps()    const { std::lock_guard g(mu_); return belief_swaps_; }
    int actor_conflicts() const { std::lock_guard g(mu_); return actor_conflicts_; }
    std::size_t live()    const { std::lock_guard g(mu_); return tracks_.size(); }

private:
    // ── Locked internals ─────────────────────────────────────────────────────
    void tick_locked(double now) {
        // The tracker's clock still bounds association (a dormant track is only
        // a candidate while it is alive), but it no longer decides death.
        now_ = now;
        reap_locked();
    }

    /// Reap against the evidence watermark when a producer of one is attached
    /// (see expect_evidence); otherwise against the tracker's clock, which is
    /// the same thing when there is only one clock.
    void reap_locked() {
        const double clock = awaits_evidence_ ? evidence_through_ : now_;
        for (auto it = tracks_.begin(); it != tracks_.end(); ) {
            const auto& ls = it->second.last_seen;
            if (ls && (clock - *ls) > cfg_.track_extinction_sec) {
                emit_locked(it->second, *ls);
                it = tracks_.erase(it);
            } else {
                ++it;
            }
        }
    }

    int create_locked(double t, const Embedding& e) {
        const int id = next_id_++;
        Track tr;
        tr.id         = id;
        tr.first_seen = t;
        tr.mean       = e;
        tr.n_obs      = 0;
        tracks_.emplace(id, std::move(tr));
        return id;
    }

    void mark_seen_locked(int id, double t, const Embedding& e) {
        auto it = tracks_.find(id);
        if (it == tracks_.end()) return;
        Track& tr = it->second;
        tr.last_seen.reset();          // back on screen; the gap is absorbed
        update_mean(tr, e);
        (void)t;
    }

    void mark_lost_locked(int id, double last_on_screen) {
        auto it = tracks_.find(id);
        if (it == tracks_.end()) return;
        it->second.last_seen = last_on_screen;
    }

    void claim_locked(Track& t, int actor) {
        // AR-015 — if another live track already owns this actor, at least one
        // is wrong: a person cannot be in two places at once. The cause is the
        // same as a belief swap — a missed camera or scene change. Detected on
        // the update that causes it via the reverse index, not by scanning.
        auto seen = owner_index_.find(actor);
        if (seen != owner_index_.end() && seen->second != t.id
            && tracks_.count(seen->second)) {
            ++actor_conflicts_;
        }
        t.actor = actor;
        owner_index_[actor] = t.id;
    }

    void split_locked(Track& t, int new_actor) {
        ++belief_swaps_;
        const double boundary = t.last_seen.value_or(t.first_seen);
        emit_locked(t, boundary);

        // The successor inherits the embedding and the belief that caused the
        // swap, and starts at the swap frame — so the two windows abut without
        // overlapping and neither blends the two people.
        Track next;
        next.id         = next_id_++;
        next.first_seen = boundary;
        next.mean       = t.mean;
        next.belief[new_actor] = t.belief[new_actor];
        next.n_obs      = 1;
        const int old_id = t.id;
        Track stash = std::move(next);
        tracks_.erase(old_id);
        const int nid = stash.id;
        tracks_.emplace(nid, std::move(stash));
        claim_locked(tracks_.at(nid), new_actor);
    }

    void emit_locked(Track& t, double end_ts) {
        if (!on_dead_) return;
        DeadTrack d;
        d.track_id     = t.id;
        d.first_seen   = t.first_seen;
        d.last_seen    = end_ts;
        d.observations  = t.n_obs;
        d.effective_obs = t.discounted_weight;
        if (t.actor) {
            d.actor_idx = *t.actor;
            d.belief    = 1.f - std::exp(t.belief[*t.actor]);
            auto oi = owner_index_.find(*t.actor);
            if (oi != owner_index_.end() && oi->second == t.id) owner_index_.erase(oi);
        }
        on_dead_(d);
    }

    /// Most-believed actor. belief holds log(1 − P), so the strongest claim is
    /// the *most negative* entry, not the largest.
    static int argmax_belief(const Track& t) {
        int   best = -1;
        float lo   = 1e30f;
        for (const auto& [a, v] : t.belief) if (v < lo) { lo = v; best = a; }
        return best;
    }

    /// Ownership expressed as a probability. Config still carries log-odds so
    /// the knob keeps its meaning across this change.
    float own_threshold() const {
        return 1.f / (1.f + std::exp(-cfg_.ownership_logodds));
    }

    static void update_mean(Track& t, const Embedding& e) {
        // Directional mean: accumulate then re-normalise to the unit sphere, so
        // cosine against it stays a plain dot product.
        double norm = 0.0;
        for (int i = 0; i < 512; ++i) {
            t.mean[i] = t.mean[i] * static_cast<float>(t.n_obs ? t.n_obs : 1) + e[i];
            norm += static_cast<double>(t.mean[i]) * t.mean[i];
        }
        norm = norm > 0 ? std::sqrt(norm) : 1.0;
        for (int i = 0; i < 512; ++i) t.mean[i] = static_cast<float>(t.mean[i] / norm);
    }

    Config                cfg_;
    EvidenceDiscounter    discounter_;
    mutable std::mutex    mu_;
    std::map<int, Track>  tracks_;
    std::map<int, int>    owner_index_;   ///< actor_idx → live track_id (AR-015)
    DeadTrackFn           on_dead_;
    int                   next_id_{0};
    double                now_{0.0};              ///< tracker's clock (association)
    double                evidence_through_{0.0}; ///< matcher's watermark (reaping)
    bool                  awaits_evidence_{false};
    int                   dropped_votes_{0};
    int                   belief_swaps_{0};
    int                   actor_conflicts_{0};
};