#pragma once
#include "types.hpp"
#include "config.hpp"

#include <functional>

#include <cmath>
#include <cstdio>
#include <iostream>
#include <limits>
#include <map>
#include <stdexcept>
#include <string>
#include <vector>

#ifdef SAE_DEBUG
#include <opencv2/imgcodecs.hpp>
#endif
#include <opencv2/core.hpp>
#include <filesystem>

// ── TrackGallery ──────────────────────────────────────────────────────────────
// Per-film gallery expansion driven by track continuity.
//
// A single uncut face track is, by construction, one physical person: the
// face_tracker links detections frame-to-frame and clears all tracks on a scene
// cut, so a track ID never spans a cut. That continuity is a same-identity label
// the baked gallery does not have. TrackGallery exploits it in two stages:
//
//   1. Diversity buffer (per track). Every frame's embedding is offered to a
//      fixed-capacity buffer for that track. When full, the member whose best
//      similarity to the *owning actor's* gallery references is HIGHEST is
//      dropped — i.e. the pose the gallery already recognises well is the least
//      informative, so the buffer is continuously biased toward the gallery-far
//      (novel-pose) embeddings on which recognition currently fails.
//
//   2. Promotion (on confirmation). A track is "owned" by actor A once ≥N frames
//      have been accepted (by the matcher's calibrated posterior) as A. On
//      confirmation the retained buffer — the hard, gallery-far poses — is
//      promoted into A's per-film annex, subject to one safety gate: the band's
//      lower bound, re-applied across the whole store (see `store_coherence`).
//
// There is exactly one threshold here, the AR-018 band, and it is a calibrated
// probability. Novelty is no longer a threshold at all — the eviction policy
// above *orders* by gallery similarity rather than cutting at a constant, and
// the band's upper bound refuses the redundant views at the door. The raw
// cosines this replaces, expand_novelty_sim and expand_track_spread_max, are
// retired under AR-024.
//
// TRACES: AR-026 | SR-001
// The annex is in-memory and discarded when the process exits, but it is NOT
// small: every owned track contributes, so it grows with cast size and film
// length. It is therefore held as a contiguous row-major matrix with a parallel
// actor index — the same flat_emb_/flat_actor_ shape the baked gallery uses —
// and the matcher hands promoted rows to the similarity engine rather than
// scanning them with a host-side loop. The deferred pass (AR-020) needs the same
// contiguous operand to score the TBI queue against in one multiply.
//
// Promoted embeddings only help SUBSEQUENT frames and later tracks of A — the
// pipeline stays streaming, no emitted output is buffered or relabelled.

struct TrackGallery {
    explicit TrackGallery(const Config& cfg)
        : enabled_(cfg.expand_gallery)
        , buffer_size_(std::max(1, cfg.expand_buffer_size))
        , band_lo_(cfg.expand_band_lo)
        , band_hi_(cfg.expand_band_hi)
        , min_anchor_frames_(std::max(1, cfg.expand_min_anchor_frames))
        , debug_dir_(cfg.expand_debug_dir)
    {
        if (!enabled_) return;
        std::cerr << "[track_gallery] per-film expansion ON"
                  << "  buffer=" << buffer_size_
                  << "  band=[" << band_lo_ << ", " << band_hi_ << "]"
                  << "  min_anchor_frames=" << min_anchor_frames_;
        if (!debug_dir_.empty()) {
            std::filesystem::create_directories(debug_dir_);
            std::cerr << "  debug_dir=" << debug_dir_;
        }
        std::cerr << "\n";
    }

    bool enabled() const { return enabled_; }

    /// TRACES: AR-026 | SR-001
    /// The annex as a contiguous row-major matrix (annex_size() × 512) plus the
    /// parallel actor index. Only ever grows, never reordered, so a row index is
    /// stable for the life of the film — which is what lets the similarity
    /// engine hold the same rows and the actor mapping stay a plain vector.
    int annex_size() const { return static_cast<int>(annex_actor_.size()); }
    const float* annex_data() const { return annex_emb_.data(); }
    const std::vector<int>& annex_actors() const { return annex_actor_; }

    /// One annex row (512 floats). The deferred pass (AR-020) scores the whole
    /// matrix at once via annex_data(); this is for inspecting a single view.
    const float* annex_row(int i) const {
        return annex_emb_.data() + static_cast<size_t>(i) * kEmbDim;
    }

    /// TRACES: AR-026 | SR-001
    /// Hand the caller every row promoted since the previous call, appending to
    /// its buffers, and return how many. The matcher pushes these into the
    /// similarity engine so the next frame's single GEMM covers the annex —
    /// draining rather than re-reading the whole matrix keeps that O(promoted),
    /// not O(annex), per frame.
    int drain_promotions(std::vector<float>& emb_out, std::vector<int>& actor_out) {
        const int pending = annex_size() - drained_;
        if (pending <= 0) return 0;

        emb_out.insert(emb_out.end(),
                       annex_emb_.begin() + static_cast<size_t>(drained_) * kEmbDim,
                       annex_emb_.end());
        actor_out.insert(actor_out.end(),
                         annex_actor_.begin() + drained_, annex_actor_.end());
        drained_ = annex_size();
        return pending;
    }

    // Offer one observed face to its track's diversity buffer.
    //   track_id     : face_tracker track (−1 = untracked, ignored)
    //   emb          : this frame's raw embedding
    //   best_actor   : actor with the highest gallery similarity for this face
    //   best_gal_sim : that similarity (best sim to best_actor's baked+annex
    //                  refs) — a raw cosine, the last one in this class: it is
    //                  calibrated on entry and only the probability is stored
    //   accepted     : true if the matcher accepted this face as best_actor
    //   crop         : aligned crop, retained only when debug dumping is on
    void observe(int track_id, const Embedding& emb,
                 int best_actor, float best_gal_sim, bool accepted,
                 const cv::Mat& crop)
    {
        if (!enabled_ || track_id < 0) return;

        TrackState& ts = tracks_[track_id];

        /// TRACES: AR-019 | SR-005
        // accepted_frames is an EVIDENCE FLOOR, not an identity decision: it
        // asks "has this track been recognised often enough to be worth
        // promoting", never "who is it". Who it is comes from the registry.
        //
        // There used to be a per-actor tally here too, and promote() fell back
        // to its plurality winner. That made two answers to "who is this track"
        // able to coexist, and the local one ignored the Bayesian accumulation
        // entirely -- weighting thirty near-identical looks the same as thirty
        // distinct ones, which is exactly what AR-025's discounting exists to
        // stop. Since promotion only fired on the local count, the fallback was
        // reachable in the live pipeline and not merely in tests: three
        // accepted frames arrive well before a posterior crosses ownership.
        if (accepted && best_actor >= 0) ts.accepted_frames++;

        insert_into_buffer(ts, emb, best_gal_sim, crop);

        // Confirm and promote once BOTH hold: the registry owns this track, and
        // enough frames have been accepted to be worth the slots. Ownership is
        // the necessary one -- without it there is no actor to promote into.
        if (!ts.promoted && ts.registry_owner >= 0 &&
            ts.accepted_frames >= min_anchor_frames_)
            promote(track_id, ts);
    }

    /// TRACES: AR-019 | SR-005
    /// Drop the buffers of tracks the registry no longer has.
    ///
    /// `alive` is the registry's own liveness test, so this annotates the track
    /// pool rather than duplicating it — the same shape as FaceTrackerFunc's
    /// prune_boxes, and for the same reason: a second opinion about which
    /// tracks exist is a second thing that can be wrong.
    ///
    /// This replaces a `forget(int)` that had NO callers, under a comment
    /// asserting "called by the matcher when it observes a cut or track
    /// disappearance". The cut half was true by another route (clear_tracks);
    /// the disappearance half was not, so a track that died quietly kept its
    /// buffer until the next cut cleared everything.
    template <typename AlivePredicate>
    void prune_dead(const AlivePredicate& alive) {
        if (!enabled_) return;
        for (auto it = tracks_.begin(); it != tracks_.end(); ) {
            if (alive(it->first)) ++it;
            else                  it = tracks_.erase(it);
        }
    }

    /// TRACES: AR-019 | SR-005
    /// The registry's verdict on who this track is. Authoritative: it comes from
    /// the Bayesian accumulation (AR-025), where the local tally counted raw
    /// accepted frames and so weighted thirty near-identical looks the same as
    /// thirty distinct ones.
    void set_owner(int track_id, int actor_idx) {
        if (track_id < 0 || actor_idx < 0) return;
        tracks_[track_id].registry_owner = actor_idx;
    }

    /// TRACES: AR-024 | SR-005
    /// Supply the calibration belonging to the active embedder.
    ///
    /// Required, not optional. The default used to be `max(0, cosine)` — a raw
    /// cosine worn as a probability, which made `expand_band_lo = 0.90` mean
    /// "cosine above 0.9" in a test and "P(same person) above 0.9" in
    /// production. Those are wildly different gates, and nothing announced the
    /// switch. `FaceTrackerFunc` already refuses to construct without a
    /// calibration for the same reason; this now matches it.
    void set_calibration(std::function<float(float)> c) {
        if (!c) throw std::invalid_argument(
            "track_gallery: a calibration is required — the admission band is "
            "expressed in probability space (AR-024)");
        calibrate_ = std::move(c);
    }

    /// Embeddings the band refused. A store that admits nothing is as wrong as
    /// one that admits everything, and neither is visible without this.
    std::size_t band_rejected() const { return rejected_; }

    // Drop every track buffer (scene cut / EOF). Mirrors face_tracker's clear.
    void clear_tracks() { tracks_.clear(); }

private:
    struct BufEntry {
        Embedding emb;
        /// P(same person) against the owning actor's refs when observed —
        /// calibrated at the door (AR-024), so the eviction ordering below is a
        /// comparison of probabilities and the struct holds no bare cosine.
        float     gal_p{0.f};
        cv::Mat   crop;           // populated only when debug_dir_ set
    };

    struct TrackState {
        std::vector<BufEntry> buf;
        int                   accepted_frames{0};
        bool                  promoted{false};
        int                   registry_owner{-1};   ///< AR-019: authoritative
    };

    /// TRACES: AR-018, AR-024 | SR-005
    /// Banded admission: an embedding joins the store only if its similarity to
    /// something already there falls **inside a band**.
    ///
    ///   above the upper bound → redundant. It is another look at a pose the
    ///     store already covers, and adding it teaches the annex nothing while
    ///     costing a slot that a novel view could have used.
    ///   below the lower bound → suspect. Within one track every face is the
    ///     same person by construction, so an embedding unlike everything else
    ///     on the track is evidence the construction failed — a track-ID
    ///     collision or a bad detection. Admitting it is how an actor's annex
    ///     gets poisoned with someone else's face.
    ///
    /// Both bounds are calibrated probabilities, never raw cosines (AR-024): a
    /// bare similarity threshold means something different for every model and
    /// every face size, and this gate has to hold across both.
    ///
    /// The first embedding is always admitted — there is nothing for it to be
    /// redundant with, and nothing to contradict it.
    bool admit(const TrackState& ts, const Embedding& emb) const {
        if (ts.buf.empty()) return true;

        float p_max = 0.f;
        for (const auto& b : ts.buf)
            p_max = std::max(p_max, calibrate_(cosine_similarity(b.emb, emb)));

        return p_max >= band_lo_ && p_max <= band_hi_;
    }

    void insert_into_buffer(TrackState& ts, const Embedding& emb,
                            float gal_sim, const cv::Mat& crop)
    {
        if (!admit(ts, emb)) { ++rejected_; return; }

        BufEntry e;
        e.emb   = emb;
        e.gal_p = calibrate_(gal_sim);
        if (!debug_dir_.empty() && !crop.empty()) e.crop = crop.clone();

        if (static_cast<int>(ts.buf.size()) < buffer_size_) {
            ts.buf.push_back(std::move(e));
            return;
        }

        // Buffer full: evict the member the gallery recognises best (highest
        // gal_p) — least informative — but only if the newcomer is at least as
        // novel. Keeping the most gallery-far views is the whole point.
        //
        // This is an *ordering*, not a threshold: there is no constant to tune,
        // and novelty-seeking lives here rather than in a cutoff. It ranks
        // probabilities, so it says the same thing across models (AR-024).
        int   worst_i = -1;
        float worst_p = e.gal_p;   // newcomer's probability is the bar to beat
        for (int i = 0; i < static_cast<int>(ts.buf.size()); ++i) {
            if (ts.buf[i].gal_p > worst_p) {
                worst_p = ts.buf[i].gal_p;
                worst_i   = i;
            }
        }
        // worst_i == −1 → every buffered view is already more novel than the
        // newcomer; drop the newcomer instead of a better sample.
        if (worst_i >= 0) ts.buf[worst_i] = std::move(e);
    }

    void promote(int track_id, TrackState& ts) {
        ts.promoted = true;   // idempotent: never promote a track twice

        const int actor = ts.registry_owner;
        if (actor < 0) return;   // unreachable: observe() gates on this

        // ── Safety gate: the band's lower bound, across the whole store ──────
        float worst = store_coherence(ts.buf);
        if (worst < band_lo_) {
            std::cerr << "[track_gallery] track " << track_id
                      << " → actor " << actor
                      << " REJECTED (worst pairwise P=" << worst
                      << " < " << band_lo_ << ", likely ID collision)\n";
            return;
        }

        int added = 0;
        for (const auto& be : ts.buf) {
            // Row-major append: the matrix stays contiguous so the matcher can
            // hand whole blocks of new rows to the GEMM path (AR-026).
            annex_emb_.insert(annex_emb_.end(), be.emb.begin(), be.emb.end());
            annex_actor_.push_back(actor);
            if (!debug_dir_.empty() && !be.crop.empty())
                dump_mugshot(track_id, actor, added, be);
            ++added;
        }

        std::cerr << "[track_gallery] track " << track_id
                  << " confirmed actor " << actor
                  << " (" << ts.accepted_frames << " accepted frames, worst "
                  << "pairwise P=" << worst << ") — promoted " << added
                  << " views; annex now " << annex_size() << "\n";
    }

    /// TRACES: AR-018, AR-024 | SR-005
    /// The store's weakest pairwise P(same person) — the band's lower bound
    /// asked of every pair, not just of the best match at the door.
    ///
    /// `admit` compares a newcomer against its *closest* existing member, so a
    /// track that drifts gradually can chain A→B→C with every step inside the
    /// band while A and C are strangers. That is precisely the shape a track-ID
    /// collision takes when two people are merged over a slow pan, so the bound
    /// is re-asked here across all pairs before anything reaches an actor's
    /// annex. Same bound, same probability space — not a second constant.
    ///
    /// A store of one has no pair to disagree; it is coherent by construction,
    /// hence 1.
    float store_coherence(const std::vector<BufEntry>& buf) const {
        float worst = std::numeric_limits<float>::max();
        for (size_t i = 0; i < buf.size(); ++i)
            for (size_t j = i + 1; j < buf.size(); ++j)
                worst = std::min(worst,
                                 calibrate_(cosine_similarity(buf[i].emb, buf[j].emb)));
        if (worst == std::numeric_limits<float>::max()) return 1.f;
        return worst;
    }

    void dump_mugshot(int track_id, int actor, int idx, const BufEntry& be) {
#ifdef SAE_DEBUG
        char name[64];
        std::snprintf(name, sizeof(name), "trk%d_actor%d_%d_p%.3f.jpg",
                      track_id, actor, idx, be.gal_p);
        cv::imwrite((std::filesystem::path(debug_dir_) / name).string(), be.crop);
#else
        (void)track_id; (void)actor; (void)idx; (void)be;
#endif
    }

    /// cosine → P(same person). The one probability space the pipeline reasons
    /// in; see gallery_calibration.hpp's same_person_probability. Never default
    /// constructed to an identity-ish stand-in — see set_calibration.
    std::function<float(float)> calibrate_;
    std::size_t             rejected_{0};   ///< admissions refused by the band

    bool                    enabled_;
    int                     buffer_size_;
    float                   band_lo_;       ///< AR-018, from cfg.expand_band_lo
    float                   band_hi_;       ///< AR-018, from cfg.expand_band_hi
    int                     min_anchor_frames_;
    std::string             debug_dir_;

    std::map<int, TrackState> tracks_;

    /// TRACES: AR-026 | SR-001
    /// Contiguous annex matrix and its parallel actor index. `drained_` marks
    /// how much of it the similarity engine already holds.
    static constexpr int kEmbDim = 512;
    std::vector<float>        annex_emb_;     ///< annex_size() × 512, row-major
    std::vector<int>          annex_actor_;   ///< actor index per annex row
    int                       drained_{0};
};