#pragma once
/// TRACES: AR-023 | SR-002
#include "types.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

#include <nlohmann/json.hpp>
#include <opencv2/core/ocl.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

// ── GalleryCalibration ────────────────────────────────────────────────────────
// Platt-style sigmoid calibration: P(match) = σ(a · similarity + b)
// where similarity = cosine similarity ∈ [-1, 1] (dot product of L2-normalised
// ArcFace embeddings).
//
// Fitted from intra-class (positive) and inter-class (negative) pairs built
// from the gallery reference embeddings.  When calibration is invalid (too few
// positive pairs), the caller should fall back to the raw threshold.
struct GalleryCalibration {
    float a{10.f};   // scale  (positive → higher similarity → higher probability)
    float b{-5.f};   // bias   (decision boundary at similarity = -b/a)
    bool  valid{false};

    // P(match | sim) using the balanced-prior calibration.
    // Pass log_prior_odds = log(p0/(1-p0)) to adjust for a real base-rate prior p0:
    //   P(match | sim, p0) = σ(a·sim + b + log(p0/(1-p0)))
    float probability(float similarity, float log_prior_odds = 0.f) const {
        float z = a * similarity + b + log_prior_odds;
        if (z >= 0.f) return 1.f / (1.f + std::exp(-z));
        float e = std::exp(z);
        return e / (1.f + e);
    }

    // Similarity at which P(match, prior) == p
    float boundary_at(float p = 0.5f, float log_prior_odds = 0.f) const {
        return (std::log(p / (1.f - p)) - b - log_prior_odds) / a;
    }
};

// Fit a logistic sigmoid to gallery pair similarities.
// Positive pairs: same actor, different reference images.
// Negative pairs: different actors (all cross-actor embedding pairs).
// Class weights balance the (typically skewed) pos/neg ratio.
// Requires ≥2 positive pairs and ≥1 negative pair.
inline GalleryCalibration calibrate_gallery(
    const std::vector<Embedding>& flat_emb,
    const std::vector<int>&       flat_actor)
{
    constexpr int   kMinEmbeddingsForPositive = 5;
    constexpr float kDedupSimThreshold        = 1.f - 1e-7f; // sim above this = duplicate
    constexpr int   kHistBins                 = 200;

    // ── Per-actor dedup, then eligibility filter ────────────────────────────
    // Drop near-identical duplicate embeddings within each actor (e.g. the
    // same source image embedded twice). Actors left with fewer than
    // kMinEmbeddingsForPositive distinct embeddings can't supply meaningful
    // same-actor (positive) pairs, but still contribute negative
    // (cross-actor) pairs for the p(unknown) side of the fit.
    int n_actors = 0;
    for (int a : flat_actor) n_actors = std::max(n_actors, a + 1);

    std::vector<std::vector<Embedding>> by_actor(n_actors);
    for (int i = 0; i < static_cast<int>(flat_emb.size()); ++i)
        by_actor[flat_actor[i]].push_back(flat_emb[i]);

    std::vector<Embedding> flat_emb_dedup;
    std::vector<int>       flat_actor_dedup;
    std::vector<bool>      actor_eligible(n_actors, false);
    int                    n_eligible = 0;

    for (int ai = 0; ai < n_actors; ++ai) {
        std::vector<Embedding> kept;
        for (const auto& e : by_actor[ai]) {
            bool dup = false;
            for (const auto& k : kept) {
                if (cosine_similarity(e, k) > kDedupSimThreshold) { dup = true; break; }
            }
            if (!dup) kept.push_back(e);
        }
        if (static_cast<int>(kept.size()) >= kMinEmbeddingsForPositive) {
            actor_eligible[ai] = true;
            ++n_eligible;
        }
        for (auto& e : kept) {
            flat_emb_dedup.push_back(e);
            flat_actor_dedup.push_back(ai);
        }
    }

    const int n = static_cast<int>(flat_emb_dedup.size());
    std::cerr << "[calibration] dedup: " << flat_emb.size() << " -> " << n
              << " embeddings  (" << n_eligible << "/" << n_actors
              << " actors have >= " << kMinEmbeddingsForPositive
              << " distinct embeddings, eligible for positive pairs)\n";

    // Pairwise cosine-similarity matrix S = E * E^T, computed via cv::gemm
    // (BLAS/SIMD on CPU) instead of a naive O(n^2 * 512) scalar loop.
    //
    // SAE_CALIB_GEMM controls the backend:
    //   "cpu"   - cv::Mat gemm only (default if OpenCL unavailable)
    //   "gpu"   - cv::UMat gemm via OpenCL, falls back to CPU if unavailable
    //   "bench" - run both and report timings + max abs diff (default)
    cv::Mat E(n, 512, CV_32F);
    for (int i = 0; i < n; ++i)
        std::memcpy(E.ptr<float>(i), flat_emb_dedup[i].data(), 512 * sizeof(float));

    std::string mode = [] {
        const char* env = std::getenv("SAE_CALIB_GEMM");
        return env ? std::string(env) : std::string("bench");
    }();

    bool have_ocl = cv::ocl::haveOpenCL();
    std::cerr << "[calibration] OpenCL available: " << (have_ocl ? "yes" : "no");
    if (have_ocl) {
        auto dev = cv::ocl::Device::getDefault();
        std::cerr << "  device=\"" << dev.name() << "\""
                  << "  type=" << (dev.type() == cv::ocl::Device::TYPE_GPU ? "GPU"
                                  : dev.type() == cv::ocl::Device::TYPE_CPU ? "CPU"
                                  : "OTHER");
    }
    std::cerr << "  (SAE_CALIB_GEMM=" << mode << ")\n";

    std::cerr << "[calibration] similarity matrix: " << n << "x512 ("
              << (n * 512LL * sizeof(float)) / (1024 * 1024) << " MB input, "
              << (n * (long long)n * sizeof(float)) / (1024 * 1024)
              << " MB output)\n" << std::flush;

    cv::Mat S;
    bool have_cpu = false, have_gpu = false;
    cv::Mat S_cpu, S_gpu_host;

    if (mode == "cpu" || mode == "bench") {
        std::cerr << "[calibration] running cv::gemm on CPU...\n" << std::flush;
        auto t0 = std::chrono::steady_clock::now();
        cv::gemm(E, E, 1.0, cv::noArray(), 0.0, S_cpu, cv::GEMM_2_T);
        auto t1 = std::chrono::steady_clock::now();
        double secs = std::chrono::duration<double>(t1 - t0).count();
        std::cerr << "[calibration] cv::gemm CPU: " << secs << "s for "
                  << n << "x" << n << " similarity matrix\n";
        have_cpu = true;
    }

    if ((mode == "gpu" || mode == "bench") && have_ocl) {
        std::cerr << "[calibration] running cv::gemm on OpenCL device...\n" << std::flush;
        cv::UMat E_gpu, S_gpu;
        E.copyTo(E_gpu);
        auto t0 = std::chrono::steady_clock::now();
        cv::gemm(E_gpu, E_gpu, 1.0, cv::noArray(), 0.0, S_gpu, cv::GEMM_2_T);
        S_gpu_host = S_gpu.getMat(cv::ACCESS_READ).clone();
        auto t1 = std::chrono::steady_clock::now();
        double secs = std::chrono::duration<double>(t1 - t0).count();
        std::cerr << "[calibration] cv::gemm OpenCL: " << secs << "s for "
                  << n << "x" << n << " similarity matrix";
        if (have_cpu) std::cerr << "  max|diff|=" << cv::norm(S_cpu, S_gpu_host, cv::NORM_INF);
        std::cerr << "\n";
        have_gpu = true;
    } else if (mode == "gpu") {
        std::cerr << "[calibration] SAE_CALIB_GEMM=gpu requested but OpenCL is unavailable\n";
    }

    if (mode == "gpu" && have_gpu)      S = S_gpu_host;
    else if (have_cpu)                  S = S_cpu;
    else if (have_gpu)                  S = S_gpu_host;
    else                                 throw std::runtime_error("calibrate_gallery: no gemm backend produced a result");

    // ── Histogram of pairwise similarities ──────────────────────────────────
    // Instead of storing one (similarity, label) sample per pair (~n^2/2
    // entries — too many for the gradient descent below), bucket pairs into
    // kHistBins bins over sim ∈ [-1, 1] and fit the sigmoid against the
    // per-bin (positive_count, negative_count).
    std::vector<double> pos_hist(kHistBins, 0.0), neg_hist(kHistBins, 0.0);
    constexpr float bin_width = 2.f / kHistBins;

    for (int i = 0; i < n; ++i) {
        const float* row = S.ptr<float>(i);
        for (int j = i + 1; j < n; ++j) {
            int bin = static_cast<int>((row[j] + 1.f) / bin_width);
            bin = std::clamp(bin, 0, kHistBins - 1);
            if (flat_actor_dedup[i] == flat_actor_dedup[j]) {
                if (actor_eligible[flat_actor_dedup[i]]) pos_hist[bin] += 1.0;
                // same actor but ineligible (< kMinEmbeddingsForPositive) — skip
            } else {
                neg_hist[bin] += 1.0;
            }
        }
    }

    double n_pos = 0.0, n_neg = 0.0;
    for (int b = 0; b < kHistBins; ++b) { n_pos += pos_hist[b]; n_neg += neg_hist[b]; }

    if (n_pos < 2 || n_neg < 1) {
        std::cerr << "[calibration] insufficient pairs (+" << n_pos
                  << "/-" << n_neg << ") — calibration skipped\n";
        return {};
    }

    // Class weights to handle pos/neg imbalance
    double total = n_pos + n_neg;
    double w_pos = total / (2.0 * n_pos);
    double w_neg = total / (2.0 * n_neg);

    std::vector<float> bin_center(kHistBins);
    for (int b = 0; b < kHistBins; ++b) bin_center[b] = -1.f + (b + 0.5f) * bin_width;

    // Gradient descent logistic regression (2 parameters: a, bias)
    float a = 10.f, bias = -5.f;
    constexpr float lr      = 0.05f;
    constexpr int   max_iter = 20000;
    constexpr float tol     = 1e-7f;

    for (int iter = 0; iter < max_iter; ++iter) {
        double da = 0.0, db = 0.0;
        for (int b = 0; b < kHistBins; ++b) {
            float x   = bin_center[b];
            float z   = a * x + bias;
            float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z))
                                   : std::exp(z) / (1.f + std::exp(z));
            if (pos_hist[b] > 0.0) {
                double err = (sig - 1.f) * w_pos * pos_hist[b];
                da += err * x;
                db += err;
            }
            if (neg_hist[b] > 0.0) {
                double err = sig * w_neg * neg_hist[b];
                da += err * x;
                db += err;
            }
        }
        da /= total;
        db /= total;
        a    -= lr * static_cast<float>(da);
        bias -= lr * static_cast<float>(db);
        if (da * da + db * db < tol * tol) break;
    }

    // Training accuracy at P=0.5 decision boundary
    double correct = 0.0;
    for (int b = 0; b < kHistBins; ++b) {
        float z   = a * bin_center[b] + bias;
        float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z))
                               : std::exp(z) / (1.f + std::exp(z));
        correct += (sig > 0.5f) ? pos_hist[b] : neg_hist[b];
    }
    double acc = 100.0 * correct / total;

    GalleryCalibration cal{a, bias, true};
    std::cerr << "[calibration] sigmoid fitted:"
              << "  a=" << a << "  b=" << bias
              << "  boundary(P=0.5)=sim" << cal.boundary_at(0.5f)
              << "  pairs=" << static_cast<long long>(total)
              << " (+" << static_cast<long long>(n_pos) << "/-" << static_cast<long long>(n_neg) << ")"
              << "  bins=" << kHistBins
              << "  train_acc=" << acc << "%\n";
    return cal;
}

// Dumps the fitted P(match | similarity) sigmoid as both a CSV table
// (similarity, p_match) and a PNG plot, to <base_path>.csv / <base_path>.png.
inline void save_calibration_curve(const GalleryCalibration& cal, const std::string& base_path) {
    if (!cal.valid) return;

    constexpr int kSamples = 200;

    std::ofstream csv(base_path + ".csv");
    if (csv.is_open()) {
        csv << "similarity,p_match\n";
        for (int i = 0; i <= kSamples; ++i) {
            float sim = -1.f + i * (2.f / kSamples);
            csv << sim << "," << cal.probability(sim) << "\n";
        }
    }

    constexpr int W = 800, H = 600, M = 50;
    cv::Mat img(H, W, CV_8UC3, cv::Scalar(255, 255, 255));
    cv::line(img, {M, H - M}, {W - M, H - M}, {0, 0, 0}, 1);  // x-axis: similarity [-1,1]
    cv::line(img, {M, M}, {M, H - M}, {0, 0, 0}, 1);          // y-axis: P(match) [0,1]

    auto to_point = [&](float sim, float p) {
        int x = M + static_cast<int>((sim + 1.f) * 0.5f * (W - 2 * M));
        int y = (H - M) - static_cast<int>(p * (H - 2 * M));
        return cv::Point(x, y);
    };

    // P=0.5 reference line and decision boundary
    float boundary = cal.boundary_at(0.5f);
    cv::line(img, to_point(-1.f, 0.5f), to_point(1.f, 0.5f), {200, 200, 200}, 1);
    if (boundary >= -1.f && boundary <= 1.f)
        cv::line(img, to_point(boundary, 0.f), to_point(boundary, 1.f), {200, 200, 200}, 1);

    cv::Point prev = to_point(-1.f, cal.probability(-1.f));
    for (int i = 1; i <= kSamples; ++i) {
        float sim = -1.f + i * (2.f / kSamples);
        cv::Point pt = to_point(sim, cal.probability(sim));
        cv::line(img, prev, pt, {255, 0, 0}, 2);
        prev = pt;
    }

    cv::imwrite(base_path + ".png", img);
    std::cerr << "[calibration] saved curve to " << base_path << ".csv / " << base_path << ".png\n";
}

// FNV-1a 64-bit hash over the gallery's reference embeddings and actor
// assignments, used to detect whether a cached calibration is still valid.
inline uint64_t hash_gallery_embeddings(
    const std::vector<Embedding>& flat_emb,
    const std::vector<int>&       flat_actor)
{
    uint64_t h = 1469598103934665603ULL;
    constexpr uint64_t prime = 1099511628211ULL;
    auto mix = [&](const void* data, size_t n) {
        const auto* p = static_cast<const unsigned char*>(data);
        for (size_t i = 0; i < n; ++i) {
            h ^= p[i];
            h *= prime;
        }
    };

    uint64_t n = flat_emb.size();
    mix(&n, sizeof(n));
    for (const auto& emb : flat_emb)
        mix(emb.data(), emb.size() * sizeof(float));
    for (int a : flat_actor)
        mix(&a, sizeof(a));
    return h;
}

// Calibrates the gallery, reusing (cached_a, cached_b, cached_valid) if
// cached_hash matches a fresh hash of the current embeddings/actor
// assignments — the O(n^2) pairwise fit only re-runs when they actually
// change. Distinct from calibrate_gallery_cached's old sidecar-JSON-file
// design: the cache now lives in the gallery HDF5 itself (ActorGallery::
// calib_*, see gallery_store.hpp), so this takes the previous values
// in-memory rather than a file path. Sets `recomputed` so the caller (which
// holds the open gallery file/struct) knows whether it needs to persist the
// refreshed values back.
inline GalleryCalibration calibrate_gallery_cached(
    const std::vector<Embedding>& flat_emb,
    const std::vector<int>&       flat_actor,
    float    cached_a,
    float    cached_b,
    bool     cached_valid,
    uint64_t cached_hash,
    const std::string& curve_base_path,
    bool&    recomputed)
{
    uint64_t hash = hash_gallery_embeddings(flat_emb, flat_actor);
    recomputed = false;

    if (cached_hash != 0 && cached_hash == hash) {
        GalleryCalibration cal{cached_a, cached_b, cached_valid};
        std::cerr << "[calibration] using cached calibration from gallery"
                  << " (a=" << cal.a << " b=" << cal.b
                  << " valid=" << cal.valid << ")\n";
        if (!curve_base_path.empty()) save_calibration_curve(cal, curve_base_path);
        return cal;
    }
    if (cached_hash != 0)
        std::cerr << "[calibration] cached calibration is stale (embeddings changed), "
                     "recomputing\n";

    auto t0 = std::chrono::steady_clock::now();
    GalleryCalibration cal = calibrate_gallery(flat_emb, flat_actor);
    auto t1 = std::chrono::steady_clock::now();
    std::cerr << "[calibration] fit took "
              << std::chrono::duration<double>(t1 - t0).count() << "s for "
              << flat_emb.size() << " embeddings\n";

    if (!curve_base_path.empty()) save_calibration_curve(cal, curve_base_path);
    recomputed = true;
    return cal;
}