#pragma once
/// TRACES: AR-005, AR-029, 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;
}

// ── crop_sharpness ────────────────────────────────────────────────────────────
/// TRACES: AR-029 | SR-002
//
// Normalised variance of the Laplacian over the aligned 112×112 crop: the AR-029
// sharpness axis. Returns -1 for an empty crop (unscored), matching the
// DetectedFace sentinel.
//
//     sharpness = Var(∇²I) / Var(I)
//
// Two normalisations, each removing a quantity that would otherwise be read as
// blur:
//
//   - **Divided by the image variance, so contrast cannot leak in.** Scaling
//     intensity by α scales the Laplacian by α too, so both variances scale by α²
//     and the ratio is unchanged. A raw Var(∇²I) — the textbook measure — instead
//     falls with exposure, so a dim scene reads as soft and a graded-up one as
//     sharp. VR-012 has to locate one knee across films whose grading differs by
//     more than their focus does; an uncalibrated measure would put the knee in a
//     different place per film, which is the AR-024 failure in another metric.
//   - **Measured on the aligned crop, so size cannot leak in.** The destination
//     frame is fixed at 112×112 (AR-002 owns size, and double-counting it here
//     would make every small face read as blurred). What the ratio reports is the
//     detail actually present in the embedder's input — so a small sharp face can
//     and does outscore a large soft one. That is the claim; it is *not* a claim
//     of invariance to source resolution, because a 40 px face warped up to 112
//     genuinely carries less detail, and hiding that would defeat the point.
//
// Frequency-domain reading of why the blur ladder is monotone: with
// Var(∇²I) = ∫|ω|⁴|F(ω)|² and Var(I) = ∫|F(ω)|², the ratio is E[|ω|⁴] under the
// image's own spectral measure. Gaussian blur multiplies that measure by
// e^{-σ²|ω|²}, concentrating it at low |ω|, so the expectation falls strictly
// with σ. It is a property of the construction, not a fitted behaviour.
//
// **Three known hazards, for VR-012 to check rather than for a threshold to
// absorb.** All are recorded here because they are properties of the measure,
// visible in the dumped distribution, and neither should be papered over by a
// correction chosen before that distribution has been looked at.
//
//   1. **Border fill.** `align_face` warps with BORDER_CONSTANT, so a face
//      crossing the frame edge brings a hard black step into the crop, and a
//      step edge is high-frequency. The normalisation blunts it — the fill
//      inflates Var(I) as well as Var(∇²I) — but does not remove it, so
//      heavily-cropped faces may read sharper than they are. The fix is either a
//      validity mask or a different border mode, and the second changes what the
//      embedder is fed (AR-011).
//
//   2. **The contrast invariance is exact in the algebra and approximate in
//      8 bits.** Scaling I by α cancels exactly; what does not cancel is the
//      quantisation floor of a stored crop, which is broadband and so lands in
//      the numerator. It matters only where there is little signal left to
//      compete with it: on the AR-029 test texture a half-contrast copy reads
//      0.9% high when sharp, 24% high at sigma 1.2 and 148% high at sigma 2.5.
//      A crop that is both **dim and soft therefore reads sharper than it is** —
//      the low corner of the axis, and the corner VR-012 must put a knee in.
//
//   3. **It reports where the energy sits, not how much there is.** A crop whose
//      energy is *already* concentrated at high frequency — dense film grain,
//      a face against foliage — loses numerator and denominator together under
//      blur, so the ratio moves less than the damage does. Measured on a
//      flat-spectrum synthetic, an anisotropic (motion) smear even makes it rise,
//      because the surviving perpendicular detail really is as fine as before.
//      Natural crops have the low-frequency mass that keeps the denominator
//      steady, and on those both ladders fall (see the AR-029 tests, which use a
//      1/f texture for exactly this reason). The same property means the axis
//      conflates focus with intrinsic texture — a bearded face outscores a smooth
//      one at equal focus — which is true of every no-reference sharpness measure
//      and is why AR-028 carries the number instead of thresholding on it.
inline float crop_sharpness(const cv::Mat& crop) {
    if (crop.empty()) return -1.f;

    cv::Mat gray;
    if (crop.channels() == 3) cv::cvtColor(crop, gray, cv::COLOR_BGR2GRAY);
    else                      gray = crop;

    cv::Mat lap;
    cv::Laplacian(gray, lap, CV_32F, 3);

    cv::Scalar mean_i, sd_i, mean_l, sd_l;
    cv::meanStdDev(gray, mean_i, sd_i);
    cv::meanStdDev(lap,  mean_l, sd_l);

    const double var_i = sd_i[0] * sd_i[0];
    // A flat crop has no detail to be sharp or soft about, and the ratio is 0/0.
    // Zero is the honest answer and keeps the axis finite; -1 would claim the
    // face was never scored, which is a different fact.
    if (var_i < 1e-6) return 0.f;

    return static_cast<float>((sd_l[0] * sd_l[0]) / var_i);
}

// ── 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;
}
