#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.
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
    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};

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

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

    struct Config {
        double extinction_sec{5.0};       ///< how long a lost track stays revivable
        float  ownership_logodds{2.0f};   ///< belief needed to own a track (~0.88 posterior)
    };

    /// 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 live 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).
        std::vector<Track*> candidates() {
            std::vector<Track*> out;
            out.reserve(reg_.tracks_.size());
            for (auto& [id, t] : reg_.tracks_) 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); }

    // ── 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;
        const float w = discounter_.weight(t.views, t.n_obs, 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_obs;

        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.
    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) {
        for (auto it = tracks_.begin(); it != tracks_.end(); ) {
            const auto& ls = it->second.last_seen;
            if (ls && (now - *ls) > cfg_.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);
    }

    static float logistic(float z) {
        return z >= 0 ? 1.f / (1.f + std::exp(-z))
                      : std::exp(z) / (1.f + std::exp(z));
    }

    static float logit(float p) {
        const float eps = 1e-6f;
        p = std::min(1.f - eps, std::max(eps, p));
        return std::log(p / (1.f - p));
    }

    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};
    int                   dropped_votes_{0};
    int                   belief_swaps_{0};
    int                   actor_conflicts_{0};
};