// sae_embed — Python module wrapping FaceEmbedderEngine (SCRFD + ArcFace).
//
// Loads both ONNX sessions once per FaceEmbedder instance, then embeds many
// images via repeated embed() calls — avoiding the per-process model-load
// cost of the embed_faces CLI when embedding a large gallery.
//
// Beyond whole-image embed(), the individual pipeline stages are exposed —
// detect(), align_face(), embed_crop() — plus the gallery calibration. A study
// that needs to step between stages (a different landmark source, a degraded
// crop) drives the shipped C++ from Python rather than re-implementing
// detection, alignment, the ArcFace warp or the Platt fit in numpy. Those
// re-implementations drift from what ships, and the calibration is the one
// that must not: AR-024 requires every similarity to pass through
// GalleryCalibration::probability, never a bare cosine.

#include "face_embedder_engine.hpp"
#include "gallery/gallery_calibration.hpp"
#include "gallery/gallery_store.hpp"

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>

#include <array>
#include <cstring>
#include <stdexcept>

namespace nb = nanobind;
using namespace nb::literals;

namespace {

using ImageArray = nb::ndarray<const uint8_t, nb::ndim<3>, nb::c_contig, nb::device::cpu>;

// numpy HxWx3 uint8 (BGR, as cv::imread yields) → cv::Mat sharing that buffer.
// The Mat is a view: it must not outlive the caller's array, so every use here
// copies or consumes it before returning.
cv::Mat as_mat(const ImageArray& a) {
    if (a.shape(2) != 3)
        throw std::invalid_argument("expected an HxWx3 uint8 BGR image");
    return cv::Mat(static_cast<int>(a.shape(0)), static_cast<int>(a.shape(1)),
                   CV_8UC3, const_cast<uint8_t*>(a.data()));
}

// cv::Mat → freshly-allocated numpy array (owns its buffer).
nb::ndarray<nb::numpy, uint8_t> mat_to_numpy(const cv::Mat& m) {
    cv::Mat c = m.isContinuous() ? m : m.clone();
    auto* buf = new uint8_t[c.total() * c.elemSize()];
    std::memcpy(buf, c.data, c.total() * c.elemSize());
    nb::capsule owner(buf, [](void* p) noexcept { delete[] static_cast<uint8_t*>(p); });
    size_t shape[3] = {static_cast<size_t>(c.rows), static_cast<size_t>(c.cols),
                       static_cast<size_t>(c.channels())};
    return nb::ndarray<nb::numpy, uint8_t>(buf, 3, shape, owner);
}

nb::ndarray<nb::numpy, float> vec_to_numpy(std::vector<float>&& v) {
    auto* buf = new float[v.size()];
    std::memcpy(buf, v.data(), v.size() * sizeof(float));
    nb::capsule owner(buf, [](void* p) noexcept { delete[] static_cast<float*>(p); });
    size_t shape[1] = {v.size()};
    return nb::ndarray<nb::numpy, float>(buf, 1, shape, owner);
}

// numpy (5,2) float32 → the landmark array align_face expects. Order is
// types.hpp:60 — [0] right-eye [1] left-eye [2] nose [3] right-mouth [4] left-mouth.
std::array<cv::Point2f, 5> as_landmarks(
    const nb::ndarray<const float, nb::shape<5, 2>, nb::c_contig, nb::device::cpu>& a) {
    std::array<cv::Point2f, 5> lm;
    for (int i = 0; i < 5; ++i) lm[i] = {a(i, 0), a(i, 1)};
    return lm;
}

}  // namespace

NB_MODULE(sae_embed, m) {
    m.doc() = "SCRFD + ArcFace face embedding, models loaded once per FaceEmbedder";

    nb::class_<FaceEmbedResult>(m, "FaceResult")
        .def_ro("ok", &FaceEmbedResult::ok)
        .def_ro("error", &FaceEmbedResult::error)
        .def_ro("confidence", &FaceEmbedResult::confidence)
        .def_prop_ro("embedding", [](const FaceEmbedResult& r) -> std::optional<std::vector<float>> {
            if (!r.ok) return std::nullopt;
            return std::vector<float>(r.embedding.begin(), r.embedding.end());
        })
        .def_prop_ro("bbox", [](const FaceEmbedResult& r) {
            return std::vector<float>{r.bbox[0], r.bbox[1], r.bbox[2], r.bbox[3]};
        })
        .def_prop_ro("landmarks", [](const FaceEmbedResult& r) {
            std::vector<float> v;
            for (const auto& p : r.landmarks) { v.push_back(p.x); v.push_back(p.y); }
            return v;
        });

    nb::class_<DetectedFace>(m, "Detection")
        .def_ro("confidence", &DetectedFace::confidence)
        .def_prop_ro("bbox", [](const DetectedFace& d) {
            return std::vector<float>{d.bbox.x, d.bbox.y, d.bbox.width, d.bbox.height};
        })
        .def_prop_ro("landmarks", [](const DetectedFace& d) {
            // (5,2): [0] right-eye [1] left-eye [2] nose [3] right-mouth [4] left-mouth
            std::vector<float> v;
            for (const auto& p : d.landmarks) { v.push_back(p.x); v.push_back(p.y); }
            return v;
        });

    nb::class_<FaceEmbedderEngine>(m, "FaceEmbedder")
        .def(nb::init<std::string, std::string, float, float, int,
                      std::string, std::string>(),
             "detector_model"_a, "arcface_model"_a,
             "conf"_a = 0.5f, "nms"_a = 0.4f, "max_side"_a = 500,
             "detector_engine"_a = "", "arcface_engine"_a = "")
        .def("embed", &FaceEmbedderEngine::embed_path, "path"_a,
             nb::call_guard<nb::gil_scoped_release>(),
             "Detect the highest-confidence face in the image, align it, and "
             "return a FaceResult with its 512-d ArcFace embedding.")
        .def("embed_mat", [](FaceEmbedderEngine& e, ImageArray img) {
                 return e.embed_mat(as_mat(img).clone());
             }, "image"_a,
             "As embed(), on an in-memory HxWx3 uint8 BGR array.")
        .def("detect", [](FaceEmbedderEngine& e, ImageArray img) {
                 return e.detect(as_mat(img));
             }, "image"_a,
             "Run the configured detector. Returns every Detection, unfiltered — "
             "min_face_px is applied downstream in face_detector_node.")
        .def("embed_crop", [](FaceEmbedderEngine& e, ImageArray crop) {
                 cv::Mat c = as_mat(crop);
                 if (c.rows != 112 || c.cols != 112)
                     throw std::invalid_argument("embed_crop expects a 112x112 aligned crop");
                 Embedding emb = e.embed_crop(c);
                 return vec_to_numpy(std::vector<float>(emb.begin(), emb.end()));
             }, "crop"_a,
             "Embed a caller-supplied 112x112 aligned BGR crop. The stage-level "
             "entry point for studies that degrade or re-align a crop themselves.")
        .def("embed_crops", [](FaceEmbedderEngine& e,
                               nb::ndarray<const uint8_t, nb::ndim<4>, nb::c_contig,
                                           nb::device::cpu> crops) {
                 if (crops.shape(1) != 112 || crops.shape(2) != 112 || crops.shape(3) != 3)
                     throw std::invalid_argument("embed_crops expects (N,112,112,3) uint8 BGR");
                 const size_t n = crops.shape(0);
                 std::vector<cv::Mat> mats;
                 mats.reserve(n);
                 for (size_t i = 0; i < n; ++i)
                     mats.emplace_back(112, 112, CV_8UC3,
                                       const_cast<uint8_t*>(crops.data()) + i * 112 * 112 * 3);
                 std::vector<Embedding> out = e.embed_crops(mats);
                 auto* buf = new float[n * 512];
                 for (size_t i = 0; i < n; ++i)
                     std::memcpy(buf + i * 512, out[i].data(), 512 * sizeof(float));
                 nb::capsule owner(buf, [](void* p) noexcept { delete[] static_cast<float*>(p); });
                 size_t shape[2] = {n, 512};
                 return nb::ndarray<nb::numpy, float>(buf, 2, shape, owner);
             }, "crops"_a,
             "Batched embed_crop: (N,112,112,3) uint8 BGR in, (N,512) float32 out. "
             "The backend batches internally, so this avoids paying per-call "
             "overhead once per crop across a large study.")
        .def_prop_ro("max_batch", [](FaceEmbedderEngine& e) { return e.max_batch(); });

    m.def("align_face", [](ImageArray img,
                           nb::ndarray<const float, nb::shape<5, 2>, nb::c_contig,
                                       nb::device::cpu> landmarks)
                        -> std::optional<nb::ndarray<nb::numpy, uint8_t>> {
              cv::Mat crop = ::align_face(as_mat(img), as_landmarks(landmarks));
              if (crop.empty()) return std::nullopt;   // degenerate fit
              return mat_to_numpy(crop);
          }, "image"_a, "landmarks"_a,
          "The ArcFace 5-point similarity transform (face_utils.hpp, AR-005). "
          "Returns a 112x112 BGR crop, or None if the affine fit is degenerate. "
          "Landmark order is types.hpp:60 — right-eye, left-eye, nose, "
          "right-mouth, left-mouth.");

    m.def("enhance_for_retry", [](ImageArray img) {
              return mat_to_numpy(::enhance_for_retry(as_mat(img)));
          }, "image"_a,
          "Border-replicate pad by 50% and CLAHE, for a detector second try.");

    // ── Calibration ──────────────────────────────────────────────────────────
    // AR-024: the pipeline reasons in one probability space. Exposed so Python
    // scores through the same sigmoid the C++ matcher uses, rather than a numpy
    // copy of it that can silently disagree.
    nb::class_<GalleryCalibration>(m, "GalleryCalibration")
        .def_ro("a", &GalleryCalibration::a)
        .def_ro("b", &GalleryCalibration::b)
        .def_ro("valid", &GalleryCalibration::valid)
        .def("probability", &GalleryCalibration::probability,
             "similarity"_a, "log_prior_odds"_a = 0.f,
             "P(match | sim) = sigma(a*sim + b + log_prior_odds). Pass "
             "log_prior_odds = log(p0/(1-p0)) for a base-rate prior p0; leave it "
             "at 0 for association (is this one person), which is what the "
             "balanced fit answers — see gallery_calibration.hpp:63.")
        .def("boundary_at", &GalleryCalibration::boundary_at,
             "p"_a = 0.5f, "log_prior_odds"_a = 0.f,
             "The similarity at which P(match) == p. Diagnostic only — decisions "
             "threshold the probability, not this.")
        .def("__repr__", [](const GalleryCalibration& c) {
            return "<GalleryCalibration a=" + std::to_string(c.a) +
                   " b=" + std::to_string(c.b) +
                   (c.valid ? " valid>" : " INVALID>");
        });

    m.def("gallery_calibration", [](const std::string& gallery_path) {
              ActorGallery g = load_gallery(gallery_path);
              if (g.calib_valid) {
                  std::cerr << "[calibration] " << gallery_path << ": cached fit"
                            << " over " << g.actors.size() << " actors\n";
                  return GalleryCalibration{g.calib_a, g.calib_b, true};
              }
              // Legacy JSON galleries carry no stored fit; compute it over the
              // whole gallery, which is the point — the calibration must come
              // from the production actor population, not a handful of people.
              std::cerr << "[calibration] " << gallery_path
                        << ": no cached fit, computing over " << g.actors.size()
                        << " actors\n";
              std::vector<Embedding> flat;
              std::vector<int>       actor;
              for (size_t a = 0; a < g.actors.size(); ++a)
                  for (const auto& e : g.actors[a].embeddings) {
                      flat.push_back(e);
                      actor.push_back(static_cast<int>(a));
                  }
              return ::calibrate_gallery(flat, actor);
          }, "gallery_path"_a,
          "The production gallery's calibration — the global fit over every "
          "actor in it. Use this to score, not a fit over a handful of people: "
          "a sigmoid fitted on a few identities saturates, so its probabilities "
          "mean nothing. Reads the cached fit stored in an HDF5 gallery, or "
          "computes it over the whole gallery for a legacy JSON one.");

    m.def("calibrate_gallery", [](nb::ndarray<const float, nb::shape<-1, 512>, nb::c_contig,
                                              nb::device::cpu> emb,
                                  std::vector<int> actor) {
              const size_t n = emb.shape(0);
              if (actor.size() != n)
                  throw std::invalid_argument("embeddings and actor ids differ in length");
              std::vector<Embedding> flat(n);
              for (size_t i = 0; i < n; ++i)
                  std::memcpy(flat[i].data(), &emb(i, 0), 512 * sizeof(float));
              return ::calibrate_gallery(flat, actor);
          }, "embeddings"_a, "actor_ids"_a,
          "Fit the Platt sigmoid from intra/inter-class pairs — the same fit the "
          "gallery build performs (gallery_calibration.hpp:85). embeddings is "
          "(N,512) L2-normalised float32; actor_ids is a length-N list of "
          "0-based actor indices.");
}
