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

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>

#include <cmath>

// ── umeyama_similarity ────────────────────────────────────────────────────────
// Closed-form least-squares similarity transform (rotation + uniform scale +
// translation, 4 DoF) mapping `src` onto `dst`, by Umeyama's solution.
//
// This is the estimator InsightFace aligns with — skimage's SimilarityTransform
// is `_umeyama(..., estimate_scale=True)` — and therefore the one that produced
// the crops ArcFace and LVFace were *trained* on. The canonical warp is part of
// the input distribution, not a free implementation choice (AR-011).
//
// Deliberately **not** `cv::estimateAffinePartial2D(..., cv::RANSAC)`:
//
//   - A robust estimator earns a small residual by discarding the points that
//     disagree with the model. On a turned face those are precisely the
//     foreshortened landmarks — the pose signal AR-030 exists to measure. RANSAC
//     would suppress exactly the quantity we want to read.
//   - With five points and a two-point minimal sample there is almost no
//     redundancy, so it cannot distinguish a mis-detected landmark from honest
//     out-of-plane rotation. The robustness is nominal.
//   - It is RNG-driven (`cv::theRNG()` is thread-local); this is exact, so
//     replay determinism stops depending on thread scheduling.
//
// Returns an empty Mat when the source points are degenerate (all coincident).
inline cv::Mat umeyama_similarity(const std::array<cv::Point2f, 5>& src,
                                  const std::array<cv::Point2f, 5>& dst) {
    constexpr int N = 5;

    double mu_sx = 0, mu_sy = 0, mu_dx = 0, mu_dy = 0;
    for (int i = 0; i < N; ++i) {
        mu_sx += src[i].x; mu_sy += src[i].y;
        mu_dx += dst[i].x; mu_dy += dst[i].y;
    }
    mu_sx /= N; mu_sy /= N; mu_dx /= N; mu_dy /= N;

    // var_src and the cross-covariance Σ = (1/N) Σ (d - μ_d)(s - μ_s)ᵀ
    double var_s = 0;
    cv::Matx22d sigma = cv::Matx22d::zeros();
    for (int i = 0; i < N; ++i) {
        const double sx = src[i].x - mu_sx, sy = src[i].y - mu_sy;
        const double dx = dst[i].x - mu_dx, dy = dst[i].y - mu_dy;
        var_s += sx * sx + sy * sy;
        sigma(0, 0) += dx * sx; sigma(0, 1) += dx * sy;
        sigma(1, 0) += dy * sx; sigma(1, 1) += dy * sy;
    }
    var_s /= N;
    sigma *= 1.0 / N;

    if (var_s < 1e-12) return {};   // every source point coincides — no scale

    cv::Mat w, u, vt;
    cv::SVD::compute(cv::Mat(sigma), w, u, vt, cv::SVD::FULL_UV);

    const cv::Matx22d U (u.at<double>(0,0),  u.at<double>(0,1),
                         u.at<double>(1,0),  u.at<double>(1,1));
    const cv::Matx22d Vt(vt.at<double>(0,0), vt.at<double>(0,1),
                         vt.at<double>(1,0), vt.at<double>(1,1));

    // A similarity may rotate but never mirror: if the fit came out
    // orientation-reversing, flip the least-significant singular direction.
    cv::Matx22d S = cv::Matx22d::eye();
    if (cv::determinant(U) * cv::determinant(Vt) < 0) S(1, 1) = -1;

    const cv::Matx22d R = U * S * Vt;
    const double c = (w.at<double>(0) * S(0,0) + w.at<double>(1) * S(1,1)) / var_s;

    cv::Mat M(2, 3, CV_64F);
    M.at<double>(0,0) = c * R(0,0); M.at<double>(0,1) = c * R(0,1);
    M.at<double>(1,0) = c * R(1,0); M.at<double>(1,1) = c * R(1,1);
    M.at<double>(0,2) = mu_dx - c * (R(0,0) * mu_sx + R(0,1) * mu_sy);
    M.at<double>(1,2) = mu_dy - c * (R(1,0) * mu_sx + R(1,1) * mu_sy);
    return M;
}

// ── Alignment ─────────────────────────────────────────────────────────────────
// The 5-point fit, plus what it could not explain.
//
// `residual` is the RMS landmark error in **canonical 112×112 pixels** after the
// best similarity fit. Two properties make it the AR-030 visibility measure:
//
//   - The similarity transform absorbs rotation, uniform scale and translation
//     exactly, so the residual is by construction the part of the deformation a
//     similarity *cannot* explain — out-of-plane rotation and foreshortening,
//     plus landmark noise. In-plane roll contributes nothing. The "roll must not
//     read as yaw" failure is excluded structurally rather than by tuning.
//   - The destination frame is fixed, so a 40 px face and a 400 px face are both
//     measured in the same canonical space. The measure cannot silently
//     re-express face size (already AR-002's job) the way a raw-pixel one would.
//
// It also responds to occlusion and to plainly broken landmark sets, which a
// yaw-angle estimator by construction does not.
struct Alignment {
    cv::Mat M;               ///< 2×3 CV_64F: source pixels → canonical 112×112
    float   residual{0.f};   ///< RMS canonical-pixel error; 0 ⇒ a perfect fit
    bool    ok{false};       ///< false ⇒ degenerate landmarks, no transform
};

/// Fit the canonical ArcFace template to `landmarks` and report the misfit.
inline Alignment estimate_alignment(const std::array<cv::Point2f, 5>& landmarks) {
    std::array<cv::Point2f, 5> dst;
    for (int i = 0; i < 5; ++i) dst[i] = {kArcFaceRef[i][0], kArcFaceRef[i][1]};

    Alignment a;
    a.M = umeyama_similarity(landmarks, dst);
    if (a.M.empty()) return a;

    double sq = 0;
    for (int i = 0; i < 5; ++i) {
        const double x = a.M.at<double>(0,0) * landmarks[i].x
                       + a.M.at<double>(0,1) * landmarks[i].y + a.M.at<double>(0,2);
        const double y = a.M.at<double>(1,0) * landmarks[i].x
                       + a.M.at<double>(1,1) * landmarks[i].y + a.M.at<double>(1,2);
        const double ex = x - dst[i].x, ey = y - dst[i].y;
        sq += ex * ex + ey * ey;
    }
    a.residual = static_cast<float>(std::sqrt(sq / 5.0));
    a.ok = true;
    return a;
}

// ── align_face ────────────────────────────────────────────────────────────────
// Produces a 112×112 BGR crop using the ArcFace 5-point similarity transform.
// Returns an empty Mat if the fit fails (degenerate detection). When
// `residual_out` is non-null it receives the AR-030 misfit for the same fit —
// free, since the transform has already been computed.
inline cv::Mat align_face(const cv::Mat& img,
                          const std::array<cv::Point2f, 5>& landmarks,
                          float* residual_out = nullptr) {
    const Alignment a = estimate_alignment(landmarks);
    if (!a.ok) return {};
    if (residual_out) *residual_out = a.residual;

    cv::Mat crop;
    cv::warpAffine(img, crop, a.M, {112, 112},
                   cv::INTER_LINEAR, cv::BORDER_CONSTANT, {0, 0, 0});
    return crop;
}

// ── enhance_for_retry ────────────────────────────────────────────────────────
// Used when initial face detection finds nothing. Pads the image by 50%
// (border-replicated, so the detector doesn't see a hard edge) and applies
// CLAHE to boost local contrast, giving the detector a second try.
inline cv::Mat enhance_for_retry(const cv::Mat& img) {
    cv::Mat padded;
    const int pad_x = img.cols / 4;
    const int pad_y = img.rows / 4;
    cv::copyMakeBorder(img, padded, pad_y, pad_y, pad_x, pad_x, cv::BORDER_REPLICATE);

    cv::Mat lab;
    cv::cvtColor(padded, lab, cv::COLOR_BGR2Lab);
    std::vector<cv::Mat> channels;
    cv::split(lab, channels);
    cv::createCLAHE(2.0, cv::Size(8, 8))->apply(channels[0], channels[0]);
    cv::merge(channels, lab);

    cv::Mat out;
    cv::cvtColor(lab, out, cv::COLOR_Lab2BGR);
    return out;
}

// ── l2_normalise ──────────────────────────────────────────────────────────────
inline Embedding l2_normalise(const float* row) {
    float norm = 0.f;
    for (int d = 0; d < 512; ++d) norm += row[d] * row[d];
    norm = std::sqrt(norm);
    if (norm < 1e-6f) norm = 1e-6f;
    Embedding emb;
    for (int d = 0; d < 512; ++d) emb[d] = row[d] / norm;
    return emb;
}