#pragma once
/// TRACES: AR-007, AR-008, AR-024 | SR-002
///
/// FaceTrackerFunc — KPN node that links face detections into tracks.
///
/// **The registry is the tracker's state.** The node owns no track map of its
/// own: it drives `TrackRegistry` through a `FrameScope` and reads the same
/// `Track` objects everything else reads. Two parallel copies could disagree,
/// and every divergence would surface as a wrong presence window rather than as
/// a crash — silently, and only in the output.
///
/// **One candidate pool** (AR-008). `last_seen` alone distinguishes a track that
/// is on screen from one that is dormant, and it only affects whether IoU means
/// anything. There is no parked pool and no revival branch: re-associating a
/// track whose face was lost — across a cut or not — is ordinary inter-frame
/// association, and it falls out of the embedding comparison already being done.
///
/// Assignment cost (track i, detection j):
///
///   p     = P(same person | cosine(track mean, detection))   ← calibrated
///   alpha = base weight, or 0 when position carries no information
///   cost  = alpha·(1 − IoU) + (1 − alpha)·(1 − p)
///
/// gated to INF unless the pair is admissible on position *or* on identity.
///
/// **alpha is frame- and track-dependent** (AR-007). It falls to 0 —
/// embedding only — when either:
///   - the frame is flagged `is_cut` / `is_scene_boundary`: the viewpoint
///     changed, so the same person is at a new position; or
///   - the track is dormant (`last_seen` set): time has passed since its box was
///     last observed, so that box is stale regardless of cuts.
/// Both are the same statement — spatial continuity is broken — arrived at from
/// two directions, which is why they collapse into one rule rather than two
/// branches.
///
/// **Everything is thresholded in probability space** (AR-024). The cosine goes
/// through the calibration before it is compared to anything; the raw-cosine
/// constants `track_max_embed_dist` and `cut_revive_sim` are retired.

#include "types.hpp"
#include "config.hpp"
#include "track_registry.hpp"

#include <algorithm>
#include <cmath>
#include <functional>
#include <iostream>
#include <limits>
#include <map>
#include <memory>
#include <stdexcept>
#include <string_view>
#include <vector>

struct FaceTrackerFunc {
    static constexpr std::string_view label() { return "face_tracker"; }

    /// cosine similarity → P(same person). Supplied by the caller so the fit
    /// belonging to the active embedder is used (AR-023/AR-024) — the same
    /// pattern, and normally the same function object, as
    /// `EvidenceDiscounter::Calibrate`.
    using Calibrate = std::function<float(float)>;

    /// The registry is a constructor argument, not an option: a tracker without
    /// one would have to keep its own tracks, which is the defect this replaces.
    FaceTrackerFunc(const Config& cfg,
                    std::shared_ptr<TrackRegistry> registry,
                    Calibrate calibrate)
        : registry_(std::move(registry))
        , calibrate_(std::move(calibrate))
        , alpha_base_(cfg.track_alpha)
        , min_iou_(cfg.track_min_iou)
        , min_assoc_prob_(cfg.track_assoc_min_prob)
    {
        if (!registry_)
            throw std::invalid_argument("face_tracker: registry must not be null");
        if (!calibrate_)
            throw std::invalid_argument("face_tracker: a calibration is required — "
                                        "association is decided in probability space");

        std::cerr << "[face_tracker] alpha_base=" << alpha_base_
                  << "  min_iou=" << min_iou_
                  << "  min_assoc_prob=" << min_assoc_prob_ << "\n";
    }

    TrackedSceneFrame operator()(EmbeddedSceneFrame ef) {
        if (ef.source.eof) {
            // Deliberately does *not* flush the registry. The identity matcher
            // runs downstream and its votes for the final frames are still in
            // flight; reaping here would drop them (they would land on ids that
            // no longer exist and show up as dropped_votes). AR-016's flush
            // belongs at the pipeline's termination point, after the last vote.
            boxes_.clear();
            TrackedSceneFrame out;
            out.source = std::move(ef.source);
            return out;
        }

        const double t     = ef.source.timestamp_sec;
        const int    n_det = static_cast<int>(ef.embeddings.size());

        // Unconditional: the clock must advance on frames with no detections
        // too, or a track only dies when some unrelated face happens to appear
        // and a film that ends mid-track never closes it (AR-013).
        auto scope = registry_->begin_frame(t);

        // One pool (AR-008) — on-screen and dormant tracks compete together.
        std::vector<Track*> cands = scope.candidates();
        const int n_trk = static_cast<int>(cands.size());

        prune_boxes(cands);
        std::vector<Spatial*> sp(n_trk);
        for (int ti = 0; ti < n_trk; ++ti)
            sp[ti] = &boxes_.try_emplace(cands[ti]->id, Spatial{{}, t, false})
                            .first->second;

        // AR-007 — the frame half of the frame-dependent weighting. Both flags
        // say the same thing to the tracker: whatever was at that position is
        // not there any more.
        const bool viewpoint_change =
            ef.source.is_cut || ef.source.is_scene_boundary;

        // ── Cost matrix [n_trk × n_det] ──────────────────────────────────────
        constexpr float INF_COST = 1e6f;
        std::vector<std::vector<float>> cost(n_trk,
            std::vector<float>(n_det, INF_COST));

        for (int ti = 0; ti < n_trk; ++ti) {
            // Spatial continuity holds only for a track that was on screen, whose
            // box we have actually observed, on a frame that did not change the
            // viewpoint. Otherwise the box is stale and IoU is noise.
            const bool spatial_meaningful =
                sp[ti]->observed && cands[ti]->on_screen() && !viewpoint_change;
            const float alpha = spatial_meaningful ? alpha_base_ : 0.f;

            for (int di = 0; di < n_det; ++di) {
                // AR-024 — the cosine is converted before it is used for
                // anything, including the gate below.
                const float p = calibrate_(
                    cosine_similarity(cands[ti]->mean, ef.embeddings[di]));
                const float iou_v = spatial_meaningful
                                    ? iou(sp[ti]->bbox, ef.faces[di].bbox) : 0.f;

                // Either signal on its own can admit a link: a face that moved a
                // little but whose embedding degraded (blur, profile turn) is
                // still linkable on position, and a face that jumped across the
                // frame is still linkable on identity. Neither ⇒ no link.
                const bool spatial_ok  = spatial_meaningful && iou_v >= min_iou_;
                const bool identity_ok = p >= min_assoc_prob_;
                if (!spatial_ok && !identity_ok) continue;

                cost[ti][di] = alpha * (1.f - iou_v) + (1.f - alpha) * (1.f - p);
            }
        }

        // ── Hungarian assignment ─────────────────────────────────────────────
        std::vector<int> assign(n_trk, -1);
        if (n_trk > 0 && n_det > 0)
            assign = hungarian(cost, n_trk, n_det);

        // ── Build output frame ───────────────────────────────────────────────
        TrackedSceneFrame out;
        out.source     = ef.source;
        out.faces      = ef.faces;
        out.crops      = ef.crops;
        out.embeddings = ef.embeddings;
        out.track_ids.assign(n_det, -1);

        std::vector<bool> det_matched(n_det, false);

        for (int ti = 0; ti < n_trk; ++ti) {
            const int di = assign[ti];
            const int id = cands[ti]->id;
            const bool valid =
                (di >= 0 && di < n_det && cost[ti][di] < INF_COST * 0.5f);

            if (!valid) {
                // Only a track that *was* on screen can become lost, and it
                // becomes lost as of its last sighting, never as of now — the
                // gap after the final sighting is never claimed (AR-013). A
                // track already dormant is left alone so its extinction clock
                // keeps running from the right instant.
                if (cands[ti]->on_screen()) scope.mark_lost(id, sp[ti]->last_ts);
                continue;
            }

            scope.mark_seen(id, t, ef.embeddings[di]);
            sp[ti]->bbox     = ef.faces[di].bbox;
            sp[ti]->last_ts  = t;
            sp[ti]->observed = true;
            det_matched[di]  = true;
            out.track_ids[di] = id;
        }

        for (int di = 0; di < n_det; ++di) {
            if (det_matched[di]) continue;
            const int id = scope.create(t, ef.embeddings[di]);
            boxes_[id] = Spatial{ef.faces[di].bbox, t, true};
            out.track_ids[di] = id;
        }

        // No reaping here: begin_frame's tick owns the extinction sweep, so
        // there is exactly one place a track can die.
        return out;
    }

private:
    // ── Spatial annotation ───────────────────────────────────────────────────
    // The one piece of per-track state the registry does not hold, because it is
    // not about presence: where the face was, and when it was last seen there.
    // Keyed by registry track id and pruned against `candidates()` every frame,
    // so it cannot outlive or contradict the registry — it annotates the pool
    // rather than duplicating it.
    struct Spatial {
        cv::Rect2f bbox{};
        double     last_ts{0.0};     ///< timestamp of the last frame this track matched
        bool       observed{false};  ///< false until a detection has been assigned
    };

    // Drop boxes for ids the registry no longer has. `candidates()` is the
    // authority on what exists; anything else is a leak (and, for a reused id,
    // would be a stale box attached to a different person).
    void prune_boxes(const std::vector<Track*>& cands) {
        if (boxes_.size() == cands.size()) return;   // common case: nothing died
        std::map<int, Spatial> kept;
        for (const Track* t : cands) {
            auto it = boxes_.find(t->id);
            if (it != boxes_.end()) kept.emplace(t->id, it->second);
        }
        boxes_.swap(kept);
    }

    // IoU of two axis-aligned bounding boxes
    static float iou(const cv::Rect2f& a, const cv::Rect2f& b) {
        float ix = std::max(0.f, std::min(a.x + a.width,  b.x + b.width)
                                - std::max(a.x, b.x));
        float iy = std::max(0.f, std::min(a.y + a.height, b.y + b.height)
                                - std::max(a.y, b.y));
        float inter = ix * iy;
        if (inter <= 0.f) return 0.f;
        return inter / (a.width * a.height + b.width * b.height - inter);
    }

    // O(n³) potential-based Hungarian algorithm (Jonker-Volgenant / Kuhn-Munkres).
    // Returns assign[row] = col (0-indexed), or -1 when row is matched to a
    // padded virtual column (i.e., unmatched).  Rectangular matrices are padded
    // to square with 0-cost virtual entries so leftover rows/cols are absorbed
    // cheaply rather than being forced onto real rows/cols.
    static std::vector<int> hungarian(
        const std::vector<std::vector<float>>& C, int nr, int nc)
    {
        const int N = std::max(nr, nc);
        constexpr float INF_VAL = 1e30f;

        // Expand to N×N, filling virtual entries with 0
        std::vector<std::vector<float>> sq(N, std::vector<float>(N, 0.f));
        for (int i = 0; i < nr; ++i)
            for (int j = 0; j < nc; ++j)
                sq[i][j] = C[i][j];

        std::vector<float> u(N + 1, 0.f), v(N + 1, 0.f);
        std::vector<int>   p(N + 1, 0),   way(N + 1, 0);

        for (int i = 1; i <= N; ++i) {
            p[0] = i;
            int j0 = 0;
            std::vector<float> minv(N + 1, INF_VAL);
            std::vector<bool>  used(N + 1, false);
            do {
                used[j0] = true;
                int i0 = p[j0], j1 = -1;
                float delta = INF_VAL;
                for (int j = 1; j <= N; ++j) {
                    if (!used[j]) {
                        float cur = sq[i0-1][j-1] - u[i0] - v[j];
                        if (cur < minv[j]) { minv[j] = cur; way[j] = j0; }
                        if (minv[j] < delta) { delta = minv[j]; j1 = j; }
                    }
                }
                for (int j = 0; j <= N; ++j) {
                    if (used[j]) { u[p[j]] += delta; v[j] -= delta; }
                    else          minv[j] -= delta;
                }
                j0 = j1;
            } while (p[j0] != 0);
            do {
                int j1 = way[j0];
                p[j0]  = p[j1];
                j0     = j1;
            } while (j0);
        }

        // p[j] = row (1-indexed) assigned to column j (1-indexed)
        std::vector<int> ans(nr, -1);
        for (int j = 1; j <= N; ++j) {
            int row = p[j] - 1;
            int col = j   - 1;
            if (row >= 0 && row < nr && col < nc)
                ans[row] = col;
            // col >= nc → virtual column → row stays unmatched (-1)
        }
        return ans;
    }

    std::shared_ptr<TrackRegistry> registry_;
    Calibrate                      calibrate_;
    std::map<int, Spatial>         boxes_;   ///< track id → where it was, when
    float alpha_base_;
    float min_iou_;
    float min_assoc_prob_;
};