#pragma once
// XGBoost scene-boundary detector — C++ inference of the shipped model
// (models/scene_boundary_xgb.json), for flood-fill presence in the live pipeline.
//
// This is a POST-EOF step (like flood-fill itself): the per-film knee threshold
// needs every peak, so boundaries can only be finalized after the whole film is
// seen. The result sink collects a per-frame RGB histogram; at EOF it calls
// boundaries() with the full (timestamp, hist) series and gets back the boundary
// timestamps to flood-snap against.
//
// The feature pipeline MUST match scripts/scene_detector/train_scene_boundary.py
// exactly (206 features): a ±WIN=3s window of per-second base features + a
// 3-value debounce clock. Base per second (29):
//   video(17): sym-delta |hist(t+k)-hist(t-k)| L1 at k=1,2,4,8; per-channel corr
//              to t-1 (3); ramp bank at H=2,4,6,8,10 on z-normed hist (5);
//              per-channel energy (3); debounce phase/decay from |delta k=1| (2)
//   audio(12): same but on the log-PSD, no corr, 1 energy   [ZERO when no audio]
// then window flatten t-3..t+3 (×7) and append clock (dt, phase, decay).
//
// Audio is not available live (the pipeline has no per-second PSD stream), so the
// audio block is fed zeros — the model was trained with audio present but it is
// weak (measured) and XGBoost tolerates a constant block; the video signal
// carries the detector. (If live audio is added later, fill the block.)

#include <xgboost/c_api.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <numeric>
#include <stdexcept>
#include <string>
#include <vector>

class XGBSceneBoundary {
public:
    // Must match kHistBins in embedding_dump_node.hpp / the training dump.
    static constexpr int kHistBins = 32;          // per channel → 96-float hist
    static constexpr int kWin = 3;                // ±WIN-second window
    static constexpr double kSigmaTau = 205.0;    // SCENE_TAU (unused at infer; kept for parity docs)
    static constexpr int    kRampScales[5] = {2, 4, 6, 8, 10};

    explicit XGBSceneBoundary(const std::string& model_path) {
        if (XGBoosterCreate(nullptr, 0, &booster_) != 0)
            throw std::runtime_error("XGBoosterCreate failed");
        if (XGBoosterLoadModel(booster_, model_path.c_str()) != 0)
            throw std::runtime_error("XGBoosterLoadModel failed: " +
                                     std::string(XGBGetLastError()));
    }
    ~XGBSceneBoundary() { if (booster_) XGBoosterFree(booster_); }
    XGBSceneBoundary(const XGBSceneBoundary&) = delete;
    XGBSceneBoundary& operator=(const XGBSceneBoundary&) = delete;

    // hist:  T rows × 96 (normalised RGB histogram per second).
    // audio: T rows × B log-PSD (from AudioLogPSD; aligned to the same seconds),
    //        or empty → the audio block is filled with its zero-input values
    //        (deltas/ramp/energy 0, but debounce phase=1/decay=exp(-1), matching
    //        the Python audio_features on a zero series).
    // Returns boundary timestamps (knee-selected).
    std::vector<double> boundaries(const std::vector<std::vector<float>>& hist,
                                   const std::vector<double>& ts,
                                   const std::vector<std::vector<float>>& audio = {}) {
        const int T = static_cast<int>(hist.size());
        if (T < 2 * kWin + 2) return {};
        auto base = build_base(hist, audio);           // [T][29]
        std::vector<float> X = window_and_clock(base, hist);
        std::vector<float> prob = predict(X, T, 206);
        return knee_boundaries(prob, ts);
    }

    static std::vector<std::vector<float>> debug_base(const std::vector<std::vector<float>>& hist,
                                                      const std::vector<std::vector<float>>& audio = {}) {
        return build_base(hist, audio);
    }

    // Predict boundaries from a precomputed [rows×cols] feature matrix (for the
    // clean parity check: same bytes both sides).
    std::vector<double> boundaries_from_features(const std::vector<float>& X, int rows,
                                                 int cols, const std::vector<double>& ts) {
        auto prob = predict(X, rows, cols);
        return knee_boundaries(prob, ts);
    }
    std::vector<float> debug_predict(const std::vector<float>& X, int r, int c) {
        return predict(X, r, c);
    }
    static std::vector<int> debug_find_peaks(const std::vector<float>& p, int d) {
        return find_peaks(p, d);
    }

    // The flat [T*206] feature matrix — exposed so TRAINING uses the exact same
    // C++ features as inference (parity by construction; no numpy re-match). The
    // Python trainer reshapes to [T,206], attaches the soft target, and fits.
    static std::vector<float> feature_matrix(const std::vector<std::vector<float>>& hist,
                                             const std::vector<std::vector<float>>& audio) {
        auto base = build_base(hist, audio);
        return window_and_clock(base, hist);
    }
    static constexpr int kNFeatures = 206;

private:
    BoosterHandle booster_{nullptr};

    // ── feature builders (exact parity with the Python) ──────────────────────

    static float l1(const std::vector<float>& a, const std::vector<float>& b) {
        float s = 0; for (size_t i = 0; i < a.size(); ++i) s += std::fabs(a[i] - b[i]);
        return s;
    }

    // z-normalise each of the 96 columns across time (matches _znorm).
    static std::vector<std::vector<float>> znorm(const std::vector<std::vector<float>>& h) {
        const int T = h.size(), D = h[0].size();
        std::vector<float> mu(D, 0), sd(D, 0);
        for (auto& r : h) for (int d = 0; d < D; ++d) mu[d] += r[d];
        for (int d = 0; d < D; ++d) mu[d] /= T;
        for (auto& r : h) for (int d = 0; d < D; ++d) sd[d] += (r[d]-mu[d])*(r[d]-mu[d]);
        for (int d = 0; d < D; ++d) sd[d] = std::sqrt(sd[d]/T) + 1e-6f;
        std::vector<std::vector<float>> z(T, std::vector<float>(D));
        for (int t = 0; t < T; ++t) for (int d = 0; d < D; ++d) z[t][d] = (h[t][d]-mu[d])/sd[d];
        return z;
    }

    // ramp bank: L2 of the antisymmetric ramp-weighted sum over ±H, per scale.
    // Matches ramp_bank() (np.convolve 'same' with reversed kernel; sign folds
    // into the L2 norm so the direct antisymmetric sum is equivalent).
    static std::vector<std::array<float,5>> ramp_bank(const std::vector<std::vector<float>>& z) {
        const int T = z.size(), D = z[0].size();
        std::vector<std::array<float,5>> out(T);
        for (int k = 0; k < 5; ++k) {
            const int H = kRampScales[k];
            for (int t = 0; t < T; ++t) {
                std::vector<double> acc(D, 0.0);
                for (int l = -H; l <= H; ++l) {
                    int idx = t + l;
                    if (idx < 0 || idx >= T) continue;
                    double w = (l == 0) ? 0.0 : (l > 0 ? 1.0 : -1.0) * (double(std::abs(l))/H);
                    for (int d = 0; d < D; ++d) acc[d] += w * z[idx][d];
                }
                double n = 0; for (double v : acc) n += v*v;
                out[t][k] = static_cast<float>(std::sqrt(n));
            }
        }
        return out;
    }

    // Generic symmetric-delta + ramp + energy + debounce feature block for one
    // modality's z-normable series `raw` (hist or PSD). Fills `out` columns
    // [off .. off+width). corr=true adds the 3 per-channel corr features (video
    // only); n_energy is 3 (video, per-channel) or 1 (audio, total).
    static void modality_block(const std::vector<std::vector<float>>& raw,
                               bool corr, int n_energy,
                               std::vector<std::vector<float>>& out, int off) {
        const int T = raw.size();
        auto z = znorm(raw);
        auto rb = ramp_bank(z);
        auto sym = [&](int t, int k)->float{
            int f = std::min(T-1, t+k), b = std::max(0, t-k);
            return l1(raw[f], raw[b]);
        };
        const int B = kHistBins;   // only used for corr (video)
        for (int t = 0; t < T; ++t) {
            int o = off;
            for (int k : {1,2,4,8}) out[t][o++] = sym(t,k);
            if (corr) {
                int tp = std::max(0, t-1);
                for (int c = 0; c < 3; ++c) {
                    double ma=0, mb=0;
                    for (int i=0;i<B;++i){ ma+=raw[t][c*B+i]; mb+=raw[tp][c*B+i]; }
                    ma/=B; mb/=B; double num=0, da=0, db=0;
                    for (int i=0;i<B;++i){ double x=raw[t][c*B+i]-ma, y=raw[tp][c*B+i]-mb;
                        num+=x*y; da+=x*x; db+=y*y; }
                    out[t][o++] = float(num/(std::sqrt(da*db)+1e-9));
                }
            }
            for (int k=0;k<5;++k) out[t][o++] = rb[t][k];
            if (n_energy == 3) {
                for (int c=0;c<3;++c){ float e=0; for(int i=0;i<B;++i) e+=raw[t][c*B+i]; out[t][o++]=e; }
            } else {
                float e=0; for (float v : raw[t]) e+=v; out[t][o++]=e;
            }
            o += 2;   // debounce filled below
        }
        // debounce from this block's delta-k1 (its first column = off)
        std::vector<float> d1(T); for (int t=0;t<T;++t) d1[t]=out[t][off];
        auto clk = debounce_phase(d1);
        // debounce sits at the end of the block: off + 4(deltas) + (corr?3:0) + 5(ramp) + n_energy
        int deb = off + 4 + (corr?3:0) + 5 + n_energy;
        for (int t=0;t<T;++t){ out[t][deb]=clk[t].first; out[t][deb+1]=clk[t].second; }
    }

    // per-second base = video(17) + audio(12). Audio empty → its block is the
    // zero-series result (deltas/ramp/energy 0, debounce phase=1/decay=exp(-1)).
    static std::vector<std::vector<float>> build_base(const std::vector<std::vector<float>>& hist,
                                                      const std::vector<std::vector<float>>& audio) {
        const int T = hist.size();
        std::vector<std::vector<float>> base(T, std::vector<float>(29, 0.0f));
        modality_block(hist, /*corr=*/true, /*n_energy=*/3, base, /*off=*/0);   // video → 0..16
        if (!audio.empty() && int(audio.size()) == T) {
            modality_block(audio, /*corr=*/false, /*n_energy=*/1, base, /*off=*/17); // audio → 17..28
        } else {
            // zero-series audio: deltas/ramp/energy already 0; only debounce differs.
            auto clk = debounce_phase(std::vector<float>(T, 0.0f));
            for (int t=0;t<T;++t){ base[t][27]=clk[t].first; base[t][28]=clk[t].second; }
        }
        return base;
    }

    // matches debounce_phase(): 90th-pct peaks, dt=time since last, phase/decay.
    static std::vector<std::pair<float,float>> debounce_phase(const std::vector<float>& sig) {
        const int T = sig.size();
        std::vector<float> s(sig); std::sort(s.begin(), s.end());
        float thr = s[std::min(T-1, int(0.90*T))];
        std::vector<std::pair<float,float>> out(T);
        int last = -1000000000;
        for (int t=0;t<T;++t){
            if (sig[t] > thr) last = t;
            double dt = (last < -100000000) ? kSigmaTau : double(t - last);
            out[t] = { float(std::min(1.0, dt/kSigmaTau)), float(std::exp(-dt/kSigmaTau)) };
        }
        return out;
    }

    // window flatten (t-3..t+3, edge-pad) + append the 3-value film clock.
    static std::vector<float> window_and_clock(const std::vector<std::vector<float>>& base,
                                               const std::vector<std::vector<float>>& hist) {
        const int T = base.size(), d = base[0].size();   // d=29
        // film-level clock: time-since-last-peak on the |delta k1| video signal
        // (base col 0), same as per_second_matrix's `clock`.
        std::vector<float> sig(T); for (int t=0;t<T;++t) sig[t]=base[t][0];
        std::vector<float> ss(sig); std::sort(ss.begin(), ss.end());
        float thr = ss[std::min(T-1, int(0.90*T))];
        std::vector<float> X; X.reserve(size_t(T)*206);
        int last=-1000000000;
        for (int t=0;t<T;++t){
            for (int off=-kWin; off<=kWin; ++off){
                int idx = std::min(T-1, std::max(0, t+off));
                for (int j=0;j<d;++j) X.push_back(base[idx][j]);
            }
            if (sig[t] > thr) last=t;
            double dt=(last<-100000000)?kSigmaTau:double(t-last);
            X.push_back(float(dt));
            X.push_back(float(std::min(1.0, dt/kSigmaTau)));
            X.push_back(float(std::exp(-dt/kSigmaTau)));
        }
        return X;
    }

    std::vector<float> predict(const std::vector<float>& X, int rows, int cols) {
        DMatrixHandle dm;
        if (XGDMatrixCreateFromMat(X.data(), rows, cols, std::nanf(""), &dm) != 0)
            throw std::runtime_error("XGDMatrixCreateFromMat failed");
        bst_ulong out_len = 0; const float* out = nullptr;
        if (XGBoosterPredict(booster_, dm, 0, 0, 0, &out_len, &out) != 0)
            throw std::runtime_error("XGBoosterPredict failed");
        std::vector<float> p(out, out + out_len);
        XGDMatrixFree(dm);
        for (auto& v : p) v = std::clamp(v, 0.f, 1.f);
        return p;
    }

    // Exact replica of scipy.signal.find_peaks(x, distance=d):
    //   1. local maxima (plateau-aware: rising then falling, midpoint of a flat top)
    //   2. keep peaks by DESCENDING height; drop any within `d` of an already-kept
    //      taller peak. This is height-priority, NOT the greedy left-to-right merge
    //      — the two give different peak sets and hence a different knee.
    static std::vector<int> find_peaks(const std::vector<float>& x, int d) {
        const int n = x.size();
        std::vector<int> mid;
        int i = 1;
        while (i < n-1) {
            if (x[i-1] < x[i]) {
                int ahead = i+1;
                while (ahead < n-1 && x[ahead] == x[i]) ahead++;
                if (x[ahead] < x[i]) mid.push_back((i + ahead - 1) / 2);
                i = ahead;
            } else i++;
        }
        // height-priority distance filter (scipy's _select_by_peak_distance)
        std::vector<int> order(mid.size());
        for (size_t k = 0; k < mid.size(); ++k) order[k] = k;
        std::sort(order.begin(), order.end(),
                  [&](int a, int b){ return x[mid[a]] < x[mid[b]]; });  // ascending
        std::vector<char> keep(mid.size(), 1);
        for (int j = int(order.size())-1; j >= 0; --j) {   // tallest first
            int k = order[j];
            if (!keep[k]) continue;
            for (int l = k-1; l >= 0 && mid[k]-mid[l] < d; --l) keep[l] = 0;
            for (int r = k+1; r < int(mid.size()) && mid[r]-mid[k] < d; ++r) keep[r] = 0;
        }
        std::vector<int> out;
        for (size_t k = 0; k < mid.size(); ++k) if (keep[k]) out.push_back(mid[k]);
        return out;
    }

    // knee threshold on peak heights → boundary timestamps (matches knee_boundaries).
    static std::vector<double> knee_boundaries(const std::vector<float>& prob,
                                               const std::vector<double>& ts,
                                               int min_gap = 5) {
        std::vector<int> pk = find_peaks(prob, min_gap);
        if (pk.size() < 5) {
            std::vector<double> r; for (int i : pk) r.push_back(ts[i]); return r;
        }
        std::vector<float> h; for (int i : pk) h.push_back(prob[i]);
        std::sort(h.begin(), h.end(), std::greater<float>());
        int n = h.size(); float h0 = h.front() + 1e-9f;
        int kbest = 0; double dmax = -1;
        for (int i = 0; i < n; ++i) {
            double x = double(i)/(n-1);
            double yv = h[i]/h0;
            double chord = (h[0]/h0) + ((h[n-1]/h0)-(h[0]/h0))*x;
            if (chord - yv > dmax) { dmax = chord - yv; kbest = i; }
        }
        float knee = h[kbest];
        std::vector<double> out;
        for (int i : pk) if (prob[i] >= knee) out.push_back(ts[i]);
        return out;
    }
};