// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "PixelRefine.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>

#include <Eigen/Dense>
#include <ceres/ceres.h>
#include <ceres/covariance.h>
#include <ceres/rotation.h>

#include "../geom_refinement/LatticeReduction.h"

namespace {

// Per-pixel observation, in *raw* detector counts (no per-pixel solid-angle or
// polarization correction - same units the "normal" integrator works in; the
// per-reflection polarization correction is applied via ReflGroup::pol).
struct PixelObs {
    double x, y;          // detector pixel coordinate
    double Iobs;          // raw pixel value (signal + background)
    double Ibkg;          // local background estimate (per-shoebox level, raw counts)
    double weight;        // 1 / sigma_pixel
};

// One reflection together with the pixels of its shoebox.
struct ReflGroup {
    int h, k, l;
    double d;
    double Itrue;             // reference intensity (held fixed)
    double R_bw_sq;           // bandwidth radial-width^2 contribution (0 = monochromatic)
    double pol;               // per-reflection polarization correction (raw = true * pol)
    double Ibkg;              // local flat background (raw counts, constant over the shoebox)
    double predicted_x, predicted_y;
    double R1_eff = 0.0;          // tangential profile width to use (Term 2; 0 => fall back to data.R[1])
    double dcx = 0.0, dcy = 0.0;  // Term 3: profile recentre shift (observed centroid - predicted)
    std::vector<PixelObs> pixels;
};

double SafeInv(double x, double fallback) {
    if (!std::isfinite(x) || std::fabs(x) < 1e-30)
        return fallback;
    return 1.0 / x;
}

// Median of a vector (in place, partially reorders it).
double MedianInPlace(std::vector<double> &v) {
    if (v.empty())
        return 0.0;
    const size_t mid = v.size() / 2;
    std::nth_element(v.begin(), v.begin() + mid, v.end());
    if (v.size() % 2 == 1)
        return v[mid];
    const double hi = v[mid];
    std::nth_element(v.begin(), v.begin() + mid - 1, v.begin() + mid);
    return 0.5 * (v[mid - 1] + hi);
}

// Mask marking the *core* (radius `radius`) of every predicted spot, so that the
// local-background sampling of one reflection never picks up a neighbouring
// reflection's signal. Same idea as BraggIntegrate2D::BuildReflectionMask.
std::vector<uint8_t> BuildSpotMask(const std::vector<Reflection> &predicted, int nrefl,
                                   size_t xpixel, size_t ypixel, int radius) {
    std::vector<uint8_t> mask(xpixel * ypixel, 0);
    const double r_sq = static_cast<double>(radius) * radius;
    for (int i = 0; i < nrefl; ++i) {
        const auto &r = predicted[i];
        const int cx = static_cast<int>(std::lround(r.predicted_x));
        const int cy = static_cast<int>(std::lround(r.predicted_y));
        const int x0 = std::max(0, cx - radius);
        const int x1 = std::min<int>(static_cast<int>(xpixel) - 1, cx + radius);
        const int y0 = std::max(0, cy - radius);
        const int y1 = std::min<int>(static_cast<int>(ypixel) - 1, cy + radius);
        for (int y = y0; y <= y1; ++y) {
            for (int x = x0; x <= x1; ++x) {
                const double dx = x - r.predicted_x;
                const double dy = y - r.predicted_y;
                if (dx * dx + dy * dy <= r_sq)
                    mask[static_cast<size_t>(xpixel) * y + x] = 1;
            }
        }
    }
    return mask;
}

// Square shoebox bounds (inclusive) around a predicted spot, clamped to the
// detector. The centre is rounded to the nearest pixel with std::lround so the
// signal box is centred identically to the spot-core mask (BuildSpotMask) and
// the local-background ring (EstimateLocalBackground), which also lround. Used by
// Run and the diagnostic renderers so all three share one shoebox definition.
struct ShoeboxBox { int min_x, max_x, min_y, max_y; };
ShoeboxBox ShoeboxBounds(double px, double py, int radius, size_t xpixel, size_t ypixel) {
    const int cx = static_cast<int>(std::lround(px));
    const int cy = static_cast<int>(std::lround(py));
    return {
        std::max(cx - radius, 0),
        std::min<int>(cx + radius, static_cast<int>(xpixel) - 1),
        std::max(cy - radius, 0),
        std::min<int>(cy + radius, static_cast<int>(ypixel) - 1)
    };
}

// Local flat background around one shoebox, in raw detector counts. Samples the
// square ring shoebox_radius < max(|dx|,|dy|) <= bkg_outer_radius centred on the
// spot, dropping pixels that belong to any spot core (spot_mask) or carry a
// masked/saturated sentinel, and returns their MEAN. (The median, used previously,
// sits below the mean for a skewed background and so biases the subtracted intensity
// positive.) Mirrors the local-background region of BraggIntegrate2D.
template<class T>
bool EstimateLocalBackground(const T *image,
                             const std::vector<uint8_t> &spot_mask,
                             size_t xpixel, size_t ypixel,
                             double cx, double cy,
                             int shoebox_radius, int bkg_outer_radius,
                             double &bkg_mean) {
    const int icx = static_cast<int>(std::lround(cx));
    const int icy = static_cast<int>(std::lround(cy));
    const int x0 = std::max(0, icx - bkg_outer_radius);
    const int x1 = std::min<int>(static_cast<int>(xpixel) - 1, icx + bkg_outer_radius);
    const int y0 = std::max(0, icy - bkg_outer_radius);
    const int y1 = std::min<int>(static_cast<int>(ypixel) - 1, icy + bkg_outer_radius);

    std::vector<double> vals;
    vals.reserve(static_cast<size_t>((x1 - x0 + 1) * (y1 - y0 + 1)));
    for (int y = y0; y <= y1; ++y) {
        for (int x = x0; x <= x1; ++x) {
            // Skip the square shoebox core: that is signal, not background.
            if (std::abs(x - icx) <= shoebox_radius && std::abs(y - icy) <= shoebox_radius)
                continue;
            const size_t np = static_cast<size_t>(xpixel) * y + x;
            if (spot_mask[np])
                continue;
            const T raw = image[np];
            if (raw == std::numeric_limits<T>::max())
                continue;
            if (std::is_signed_v<T> && raw == std::numeric_limits<T>::min())
                continue;
            vals.push_back(static_cast<double>(raw));
        }
    }

    if (vals.size() < 5)
        return false;

    // Mean background, NOT median. The median of a right-skewed (Poisson) background sits
    // below the mean, so subtracting it under-subtracts and biases every weak integrated
    // intensity positive - which averages up over multiplicity into fake <I/sig> in the
    // no-signal high-resolution shells. The mean is unbiased; modern fast detectors put
    // few zingers in the ring, so it is not robustified here.
    double sum = 0.0;
    for (const double v : vals)
        sum += v;
    bkg_mean = sum / static_cast<double>(vals.size());
    return true;
}

// Per-pixel: map a detector pixel through the current geometry into the
// reference reciprocal frame. Cheap (a few trig + one rotation); depends on the
// pixel and the detector geometry, not on the lattice.
template<typename T>
void ObservedRecip(const T *beam, const T *distance_mm, const T *detector_rot,
                   double obs_x, double obs_y,
                   double pixel_size, double inv_lambda,
                   Eigen::Matrix<T, 3, 1> &e_obs_recip) {
    // PyFAI convention (left-handed for rot1/rot2): rot3 = 0 assumed.
    const T c1 = ceres::cos(detector_rot[0]);
    const T s1 = ceres::sin(detector_rot[0]);
    const T c2 = ceres::cos(detector_rot[1]);
    const T s2 = ceres::sin(detector_rot[1]);

    const T det_x = (T(obs_x) - beam[0]) * T(pixel_size);
    const T det_y = (T(obs_y) - beam[1]) * T(pixel_size);
    const T det_z = T(distance_mm[0]);

    const T t1_x = c1 * det_x + s1 * det_z;
    const T t1_y = det_y;
    const T t1_z = -s1 * det_x + c1 * det_z;

    const T x = t1_x;
    const T y = c2 * t1_y + s2 * t1_z;
    const T z = -s2 * t1_y + c2 * t1_z;

    const T inv_norm = T(1) / ceres::sqrt(x * x + y * y + z * z);

    T recip_raw[3] = {
        x * inv_norm * T(inv_lambda),
        y * inv_norm * T(inv_lambda),
        (z * inv_norm - T(1.0)) * T(inv_lambda)
    };
    e_obs_recip = Eigen::Matrix<T, 3, 1>(recip_raw[0], recip_raw[1], recip_raw[2]);
}

// Per-reflection: predicted node g_hkl, |g_hkl|^2, and the Ewald-sphere normal.
// This is the expensive part (symmetry-aware B matrix, three rotations, cross
// products) - it depends only on the lattice (p0,p1,p2) and hkl, so for a whole
// shoebox it can be computed once. Convention identical to XtalOptimizer.
template<typename T>
bool PredictedNode(const T *p0, const T *p1, const T *p2,
                   double exp_h, double exp_k, double exp_l,
                   gemmi::CrystalSystem symmetry, double inv_lambda,
                   Eigen::Matrix<T, 3, 1> &e_pred_recip,
                   Eigen::Matrix<T, 3, 1> &n_radial, T &q_sq) {
    Eigen::Matrix<T, 3, 1> e_uc_len = Eigen::Matrix<T, 3, 1>::Zero();
    Eigen::Matrix<T, 3, 3> Bmat = Eigen::Matrix<T, 3, 3>::Identity();

    if (symmetry == gemmi::CrystalSystem::Hexagonal) {
        e_uc_len << p1[0], p1[0], p1[2];
        Bmat(0, 1) = T(-0.5);
        Bmat(1, 1) = T(sqrt(3.0) / 2.0);
    } else if (symmetry == gemmi::CrystalSystem::Orthorhombic) {
        e_uc_len << p1[0], p1[1], p1[2];
    } else if (symmetry == gemmi::CrystalSystem::Tetragonal) {
        e_uc_len << p1[0], p1[0], p1[2];
    } else if (symmetry == gemmi::CrystalSystem::Cubic) {
        e_uc_len << p1[0], p1[0], p1[0];
    } else if (symmetry == gemmi::CrystalSystem::Monoclinic) {
        e_uc_len << p1[0], p1[1], p1[2];
        Bmat(0, 2) = ceres::cos(p2[0]);
        Bmat(2, 2) = ceres::sin(p2[0]);
    } else {
        const T ca = ceres::cos(p2[0]);
        const T cb = ceres::cos(p2[1]);
        const T cg = ceres::cos(p2[2]);
        const T sg = ceres::sin(p2[2]);

        e_uc_len << p1[0], p1[1], p1[2];

        Bmat(0, 1) = cg;
        Bmat(1, 1) = sg;

        const T cx = cb;
        const T cy = (ca - cb * cg) / sg;
        const T v = T(1) - cx * cx - cy * cy;
        const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0);

        Bmat(0, 2) = cx;
        Bmat(1, 2) = cy;
        Bmat(2, 2) = cz;
    }

    const T L0 = e_uc_len[0];
    const T L1 = e_uc_len[1];
    const T L2 = e_uc_len[2];

    T col0_unrot[3] = {Bmat(0, 0) * L0, Bmat(1, 0) * L0, Bmat(2, 0) * L0};
    T col1_unrot[3] = {Bmat(0, 1) * L1, Bmat(1, 1) * L1, Bmat(2, 1) * L1};
    T col2_unrot[3] = {Bmat(0, 2) * L2, Bmat(1, 2) * L2, Bmat(2, 2) * L2};

    T col0_rot[3], col1_rot[3], col2_rot[3];
    ceres::AngleAxisRotatePoint(p0, col0_unrot, col0_rot);
    ceres::AngleAxisRotatePoint(p0, col1_unrot, col1_rot);
    ceres::AngleAxisRotatePoint(p0, col2_unrot, col2_rot);

    const Eigen::Matrix<T, 3, 1> A(col0_rot[0], col0_rot[1], col0_rot[2]);
    const Eigen::Matrix<T, 3, 1> Bv(col1_rot[0], col1_rot[1], col1_rot[2]);
    const Eigen::Matrix<T, 3, 1> C(col2_rot[0], col2_rot[1], col2_rot[2]);

    const Eigen::Matrix<T, 3, 1> BxC = Bv.cross(C);
    const Eigen::Matrix<T, 3, 1> CxA = C.cross(A);
    const Eigen::Matrix<T, 3, 1> AxB = A.cross(Bv);

    const T Vol = A.dot(BxC);
    if (ceres::abs(Vol) < T(1e-12))
        return false;
    const T invV = T(1) / Vol;

    e_pred_recip = (BxC * T(exp_h) + CxA * T(exp_k) + AxB * T(exp_l)) * invV;
    q_sq = e_pred_recip.squaredNorm();

    // Ewald sphere centre at -k_i = (0,0,-inv_lambda); radial normal at g_hkl.
    const Eigen::Matrix<T, 3, 1> S_pred(
        e_pred_recip[0],
        e_pred_recip[1],
        e_pred_recip[2] + T(inv_lambda));
    const T S_pred_norm = S_pred.norm();
    if (S_pred_norm < T(1e-10))
        return false;

    n_radial = S_pred / S_pred_norm;
    return true;
}

// Pulls a scalar parameter towards an expected value with a fixed weight (the
// data-scaled prior). Identical in spirit to ScaleOnTheFly's regularizer: it is what
// keeps the per-image scale G from wandering on weakly-constrained images and
// scrambling the cross-image merge.
struct ScalarRegularizer {
    ScalarRegularizer(double weight, double expected) : weight(weight), expected(expected) {}
    template<typename T>
    bool operator()(const T *p, T *residual) const {
        residual[0] = T(weight) * (p[0] - T(expected));
        return true;
    }
    double weight;
    double expected;
};

// Anchors the orientation (angle-axis vector) to its pre-refinement value with a
// data-scaled weight. Without it the three orientation DOF chase the sparse signal
// (and the few noisy background pixels) and the per-image intensities collapse;
// with it the fit can only make a small, signal-supported sub-spot correction - the
// push that brings slightly-misaligned high-resolution reflections onto their
// shoeboxes. Mirrors the G/B regularizers in ScaleOnTheFly.
struct OrientationRegularizer {
    OrientationRegularizer(double weight, const double prior[3]) : weight(weight) {
        for (int i = 0; i < 3; ++i)
            prior_[i] = prior[i];
    }
    template<typename T>
    bool operator()(const T *p0, T *residual) const {
        for (int i = 0; i < 3; ++i)
            residual[i] = T(weight) * (p0[i] - T(prior_[i]));
        return true;
    }
    double weight;
    double prior_[3];
};

} // namespace

// ---------------------------------------------------------------------------
// Cost functor
//
//   I_pred(pixel) = G * Itrue * B_term * P_radial * P_tangential * pol + I_bkg
//
//   B_term       = exp(-B |q|^2 / 4)                    (Debye-Waller)
//   P_radial     = exp(-eps_r^2 / R0_eff^2)             (partiality: fraction of
//                                                        the mosaic blob on the
//                                                        Ewald sphere; <= 1)
//   P_tangential = exp(-eps_t^2/R1^2) / (pi R1^2)       (Gaussian spatial profile
//                                                        in the Ewald tangent plane)
//   pol          = per-reflection polarization correction (raw = true * pol),
//                  evaluated once at the predicted spot position (as in
//                  BraggIntegrate2D). 1 if polarization is disabled.
//
// Everything is in *raw* detector counts: there is no per-pixel solid-angle or
// area (Lorentz/Jacobian) weighting - each pixel counts equally, like the normal
// integrator. The tangential factor is what makes this "profile fitting"; the
// 1/(pi R1^2) normalization keeps the profile width R1 from soaking up the
// overall scale G.
//
// X-ray bandwidth: a spread in lambda is a spread in the Ewald-sphere radius,
// i.e. a purely *radial* thickening of the shell. It adds (in quadrature) a
// resolution-dependent term to the radial width:
//     R0_eff^2 = R0^2 + R_bw^2 ,   R_bw^2 = (b*lambda)^2 / (2 d^4)
// where b = relative bandwidth (sigma of dlambda/lambda). R_bw grows like 1/d^2,
// so bandwidth leaves low-resolution spots sharp and smears high-resolution ones
// radially - the pink-beam/DMM signature. R_bw_sq is a fixed per-reflection
// constant (b is known), so R0 keeps meaning "intrinsic" width (mosaic +
// divergence + beam). b = 0 makes R_bw = 0: a monochromatic no-op.
// ---------------------------------------------------------------------------
struct PixelResidual {
    PixelResidual(const PixelObs &obs, double Itrue,
                  double lambda, double pixel_size,
                  double exp_h, double exp_k, double exp_l,
                  double R_bw_sq, double pol,
                  gemmi::CrystalSystem symmetry)
        : Itrue(Itrue), Iobs(obs.Iobs), Ibkg(obs.Ibkg), weight(obs.weight),
          obs_x(obs.x), obs_y(obs.y),
          inv_lambda(1.0 / lambda), pixel_size(pixel_size),
          exp_h(exp_h), exp_k(exp_k), exp_l(exp_l),
          R_bw_sq(R_bw_sq), pol(pol), symmetry(symmetry) {
        if (std::fabs(lambda) < 1e-6)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Lambda cannot be close to zero");
    }

    // Maps a detector pixel through the current geometry + lattice into the
    // reference reciprocal frame and returns:
    //   q_sq        = |g_hkl|^2                 (predicted node, for B-factor)
    //   eps_radial  = deviation along Ewald normal (partiality direction)
    //   eps_tang_sq = squared deviation in the detector-tangential plane (profile)
    template<typename T>
    bool GeometryTerms(const T *const beam,
                       const T *const distance_mm,
                       const T *const detector_rot,
                       const T *const p0,
                       const T *const p1,
                       const T *const p2,
                       T &q_sq, T &eps_radial, T &eps_tang_sq) const {
        Eigen::Matrix<T, 3, 1> e_obs_recip;
        ObservedRecip(beam, distance_mm, detector_rot,
                      obs_x, obs_y, pixel_size, inv_lambda,
                      e_obs_recip);

        Eigen::Matrix<T, 3, 1> e_pred_recip, n_radial;
        if (!PredictedNode(p0, p1, p2, exp_h, exp_k, exp_l, symmetry, inv_lambda,
                           e_pred_recip, n_radial, q_sq))
            return false;

        const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
        eps_radial = delta_q.dot(n_radial);
        eps_tang_sq = (delta_q - eps_radial * n_radial).squaredNorm();
        return true;
    }

    // Assembles the full model intensity for the pixel from the geometry terms.
    template<typename T>
    bool Model(const T *const beam, const T *const distance_mm,
               const T *const detector_rot,
               const T *const p0, const T *const p1, const T *const p2,
               const T *const scale_factor, const T *const B, const T *const R,
               T &Ipred) const {
        T q_sq, eps_radial, eps_tang_sq;
        if (!GeometryTerms(beam, distance_mm, detector_rot,
                           p0, p1, p2, q_sq, eps_radial, eps_tang_sq))
            return false;

        if (R[0] < T(1e-10) || R[1] < T(1e-10))
            return false;

        const T B_term = ceres::exp(-B[0] * q_sq / T(4.0));

        // Separable Gaussian spot model:
        //   radial  P_r(e) = exp(-e^2/R0_eff^2)                      (peak-normalized, in (0,1])
        //   tangent g_t(e) = exp(-|e|^2/R1^2) / (pi R1^2)            [1/A^-2]
        // Every pixel counts equally (no area/Lorentz weighting); the radial factor
        // is the still-image partiality (how far the reflection sits from the Ewald
        // sphere); the overall scale is carried by the free G.
        //
        // IMPORTANT: the radial factor MUST use the same convention here as the
        // extraction's `partiality` (peak-normalized), otherwise image_scale_corr
        // = 1/(partiality*G*B) does not invert the model and a leftover, R0_eff-
        // dependent (hence resolution-dependent) factor biases the intensities.
        // R0_eff folds in the energy-bandwidth broadening via R_bw_sq.
        const T R0_eff_sq = R[0] * R[0] + T(R_bw_sq);
        const T P_radial = ceres::exp(-eps_radial * eps_radial / R0_eff_sq);
        const T P_tang = ceres::exp(-eps_tang_sq / (R[1] * R[1]))
                         / (T(M_PI) * R[1] * R[1]);

        const T signal = scale_factor[0] * T(Itrue) * B_term * P_radial * P_tang * T(pol);
        Ipred = signal + T(Ibkg);
        return true;
    }

    template<typename T>
    bool operator()(const T *const beam,
                    const T *const distance_mm,
                    const T *const detector_rot,
                    const T *const p0,
                    const T *const p1,
                    const T *const p2,
                    const T *const scale_factor,
                    const T *const B,
                    const T *const R,
                    T *residual) const {
        T Ipred;
        if (!Model(beam, distance_mm, detector_rot, p0, p1, p2, scale_factor, B, R, Ipred))
            return false;

        residual[0] = (Ipred - T(Iobs)) * T(weight);
        return true;
    }

    const double Itrue, Iobs, Ibkg, weight;
    const double obs_x, obs_y;
    const double inv_lambda;
    const double pixel_size;
    const double exp_h, exp_k, exp_l;
    const double R_bw_sq;   // bandwidth radial-width^2 contribution (0 = monochromatic)
    const double pol;       // per-reflection polarization correction
    gemmi::CrystalSystem symmetry;
};

// ---------------------------------------------------------------------------
// Per-shoebox cost functor
//
// One residual block per reflection emitting N residuals (one per shoebox pixel).
// The expensive per-reflection geometry (PredictedNode: symmetry-aware B matrix,
// three rotations, cross products) is computed ONCE; only the cheap per-pixel
// ObservedRecip + Gaussian profile run in the pixel loop. This is identical in
// value to the old one-block-per-pixel formulation but ~(pixels-per-shoebox)x
// fewer evaluations of the costly node computation. Uses the same shared helpers
// (and hence the same conventions) as PixelResidual.
// ---------------------------------------------------------------------------
struct ShoeboxResidual {
    ShoeboxResidual(const ReflGroup &g, double lambda, double pixel_size,
                    gemmi::CrystalSystem symmetry)
        : pixels(g.pixels), Itrue(g.Itrue), R_bw_sq(g.R_bw_sq), pol(g.pol),
          exp_h(g.h), exp_k(g.k), exp_l(g.l),
          inv_lambda(1.0 / lambda), pixel_size(pixel_size),
          symmetry(symmetry) {}

    template<typename T>
    bool operator()(const T *const *params, T *residual) const {
        // Parameter blocks (order matches AddParameterBlock in Run):
        // 0 beam[2] 1 distance[1] 2 detector_rot[2]
        // 3 p0[3] 4 p1[3] 5 p2[3] 6 scale[1] 7 B[1] 8 R[2]
        const T *beam = params[0];
        const T *distance_mm = params[1];
        const T *detector_rot = params[2];
        const T *p0 = params[3];
        const T *p1 = params[4];
        const T *p2 = params[5];
        const T *scale_factor = params[6];
        const T *B = params[7];
        const T *R = params[8];

        if (R[0] < T(1e-10) || R[1] < T(1e-10))
            return false;

        // --- per-reflection: computed once ---------------------------------
        Eigen::Matrix<T, 3, 1> e_pred_recip, n_radial;
        T q_sq;
        if (!PredictedNode(p0, p1, p2, exp_h, exp_k, exp_l, symmetry, inv_lambda,
                           e_pred_recip, n_radial, q_sq))
            return false;

        const T B_term = ceres::exp(-B[0] * q_sq / T(4.0));
        const T R0_eff_sq = R[0] * R[0] + T(R_bw_sq);

        // --- per-pixel loop -------------------------------------------------
        for (size_t i = 0; i < pixels.size(); ++i) {
            const PixelObs &obs = pixels[i];

            Eigen::Matrix<T, 3, 1> e_obs_recip;
            ObservedRecip(beam, distance_mm, detector_rot,
                          obs.x, obs.y, pixel_size, inv_lambda,  e_obs_recip);

            const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
            const T eps_radial = delta_q.dot(n_radial);
            const T eps_tang_sq = (delta_q - eps_radial * n_radial).squaredNorm();

            const T P_radial = ceres::exp(-eps_radial * eps_radial / R0_eff_sq);
            const T P_tang = ceres::exp(-eps_tang_sq / (R[1] * R[1]))
                             / (T(M_PI) * R[1] * R[1]);

            const T signal = scale_factor[0] * T(Itrue) * B_term * P_radial * P_tang * T(pol);
            const T Ipred = signal + T(obs.Ibkg);
            residual[i] = (Ipred - T(obs.Iobs)) * T(obs.weight);
        }
        return true;
    }

    std::vector<PixelObs> pixels;
    const double Itrue, R_bw_sq, pol;
    const double exp_h, exp_k, exp_l;
    const double inv_lambda, pixel_size;
    gemmi::CrystalSystem symmetry;
};

PixelRefine::PixelRefine(const DiffractionExperiment &experiment,
                         const std::vector<MergedReflection> &reference)
    : xpixel(experiment.GetXPixelsNum()),
      ypixel(experiment.GetYPixelsNum()),
      experiment(experiment),
      hkl_key_generator(experiment.GetScalingSettings().GetMergeFriedel(),
                        experiment.GetSpaceGroupNumber().value_or(1)) {
    for (const auto &ref: reference)
        reference_data[hkl_key_generator(ref)] = ref.I;
}

void PixelRefine::BuildParameterBlocks(const PixelRefineData &data,
                                       double beam[2], double &dist_mm,
                                       double detector_rot[2],
                                       double latt_vec0[3], double latt_vec1[3], double latt_vec2[3]) const {
    beam[0] = data.geom.GetBeamX_pxl();
    beam[1] = data.geom.GetBeamY_pxl();
    dist_mm = data.geom.GetDetectorDistance_mm();
    detector_rot[0] = data.geom.GetPoniRot1_rad();
    detector_rot[1] = data.geom.GetPoniRot2_rad();

    for (int i = 0; i < 3; ++i)
        latt_vec0[i] = latt_vec1[i] = latt_vec2[i] = 0.0;

    double beta = data.latt.GetUnitCell().beta;
    switch (data.crystal_system) {
        case gemmi::CrystalSystem::Orthorhombic:
            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
            break;
        case gemmi::CrystalSystem::Tetragonal:
            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
            latt_vec1[0] = (latt_vec1[0] + latt_vec1[1]) / 2.0;
            break;
        case gemmi::CrystalSystem::Cubic:
            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
            latt_vec1[0] = (latt_vec1[0] + latt_vec1[1] + latt_vec1[2]) / 3.0;
            break;
        case gemmi::CrystalSystem::Hexagonal:
            LatticeToRodriguesAndLengths_Hex(data.latt, latt_vec0, latt_vec1);
            break;
        case gemmi::CrystalSystem::Monoclinic:
            LatticeToRodriguesLengthsBeta_Mono(data.latt, latt_vec0, latt_vec1, beta);
            latt_vec2[0] = beta;
            break;
        default: {
            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
            const auto uc = data.latt.GetUnitCell();
            latt_vec2[0] = uc.alpha * M_PI / 180.0;
            latt_vec2[1] = uc.beta * M_PI / 180.0;
            latt_vec2[2] = uc.gamma * M_PI / 180.0;
            break;
        }
    }
}

template<class T>
void PixelRefine::SweepOrientationCell(const T *image, BraggPrediction &prediction,
                                       PixelRefineData &data) const {
    const int radius = data.shoebox_radius;
    const double beam_x = data.geom.GetBeamX_pxl();
    const double beam_y = data.geom.GetBeamY_pxl();
    const auto qnan = std::numeric_limits<double>::quiet_NaN();

    // Box-sum minus local (perimeter) background, raw counts. NaN if the box runs
    // off the detector or hits a masked/saturated pixel.
    auto integrate = [&](double px, double py) -> double {
        const int cx = static_cast<int>(std::lround(px));
        const int cy = static_cast<int>(std::lround(py));
        const int outer = radius + 1;
        if (cx - outer < 0 || cy - outer < 0 ||
            cx + outer >= static_cast<int>(xpixel) || cy + outer >= static_cast<int>(ypixel))
            return qnan;
        double sig = 0.0;
        int nsig = 0;
        std::vector<double> ring;
        ring.reserve((2 * outer + 1) * (2 * outer + 1));
        for (int y = cy - outer; y <= cy + outer; ++y) {
            for (int x = cx - outer; x <= cx + outer; ++x) {
                const T raw = image[static_cast<size_t>(xpixel) * y + x];
                if (raw == std::numeric_limits<T>::max())
                    return qnan;
                if (std::is_signed_v<T> && raw == std::numeric_limits<T>::min())
                    return qnan;
                const double v = static_cast<double>(raw);
                if (std::abs(x - cx) <= radius && std::abs(y - cy) <= radius) {
                    sig += v;
                    ++nsig;
                } else {
                    ring.push_back(v);
                }
            }
        }
        if (ring.size() < 5)
            return qnan;
        return sig - nsig * MedianInPlace(ring);
    };

    // Predict (wide band) and collect every reflection that has a reference value,
    // with its detector radius. The full set is scored - the strong low-res spots
    // anchor the CC, the weak high-res spots are what "appear" at the right cell.
    DiffractionExperiment exp_iter = experiment;
    exp_iter.BeamX_pxl(beam_x).BeamY_pxl(beam_y)
            .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
            .PoniRot1_rad(data.geom.GetPoniRot1_rad())
            .PoniRot2_rad(data.geom.GetPoniRot2_rad());
    const BraggPredictionSettings settings{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = static_cast<float>(data.ewald_dist_cutoff),
        .max_hkl = 100,
        .centering = data.centering,
        .bandwidth_sigma = static_cast<float>(data.bandwidth)
    };
    const int nrefl = prediction.Calc(exp_iter, data.latt, settings);
    const auto &predicted = prediction.GetReflections();

    struct Matched { int h, k, l; double refI; };
    std::vector<Matched> matched;
    double r_max = 0.0, r_min = std::numeric_limits<double>::max();
    for (int i = 0; i < nrefl; ++i) {
        const auto &r = predicted[i];
        const auto it = reference_data.find(hkl_key_generator(r));
        if (it == reference_data.end())
            continue;
        matched.push_back({r.h, r.k, r.l, it->second});
        const double dx = r.predicted_x - beam_x;
        const double dy = r.predicted_y - beam_y;
        const double rad = std::sqrt(dx * dx + dy * dy);
        r_max = std::max(r_max, rad);
        r_min = std::min(r_min, rad);
    }
    if (matched.size() < 20 || r_min <= 1.0 || r_max <= r_min)
        return; // too little to anchor a meaningful sweep

    // CC of the box-summed intensities against the reference, over all matched hkls.
    auto score = [&](const CrystalLattice &L) -> double {
        const Coord A = L.Astar(), B = L.Bstar(), C = L.Cstar();
        double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
        int n = 0;
        for (const auto &m : matched) {
            const Coord g = A * static_cast<float>(m.h) + B * static_cast<float>(m.k)
                            + C * static_cast<float>(m.l);
            const auto [x, y] = data.geom.RecipToDetector(g);
            if (!std::isfinite(x) || !std::isfinite(y))
                continue;
            const double I = integrate(x, y);
            if (!std::isfinite(I))
                continue;
            const double yv = m.refI;
            sx += I; sy += yv; sxx += I * I; syy += yv * yv; sxy += I * yv; ++n;
        }
        if (n < 10)
            return -2.0;
        const double nd = n;
        const double cov = sxy - sx * sy / nd;
        const double vx = sxx - sx * sx / nd;
        const double vy = syy - sy * sy / nd;
        if (!(vx > 0.0 && vy > 0.0))
            return -2.0;
        return cov / std::sqrt(vx * vy);
    };

    // Step = 1 px at the highest resolution. Range = assumed orientation/cell-scale
    // uncertainty (a few px at high res), NOT the low-res 2 px cap: the latter is
    // ~2*r_max/r_min px at high res - far too permissive, and lets the per-image CC
    // overfit. Here the low-res spots barely move (stay anchored).
    const double step = 1.0 / r_max;
    const int n_rot = std::clamp(
        static_cast<int>(std::lround(data.sweep_max_deg * M_PI / 180.0 * r_max)), 1, 25);
    const int n_scale = std::clamp(
        static_cast<int>(std::lround(data.sweep_max_cell_frac * r_max)), 1, 25);
    const Coord axes[3] = {Coord(1, 0, 0), Coord(0, 1, 0), Coord(0, 0, 1)};

    CrystalLattice best = data.latt;
    double best_cc = score(best);

    for (int round = 0; round < 2; ++round) {
        for (const auto &axis : axes) {
            CrystalLattice axis_best = best;
            double axis_cc = best_cc;
            for (int i = -n_rot; i <= n_rot; ++i) {
                if (i == 0)
                    continue;
                CrystalLattice cand = best.Multiply(RotMatrix(static_cast<float>(i * step), axis));
                const double cc = score(cand);
                if (cc > axis_cc) {
                    axis_cc = cc;
                    axis_best = cand;
                }
            }
            best = axis_best;
            best_cc = axis_cc;
        }
        CrystalLattice scale_best = best;
        double scale_cc = best_cc;
        for (int i = -n_scale; i <= n_scale; ++i) {
            if (i == 0)
                continue;
            const double s = 1.0 / (1.0 + i * step); // cell scale (1+eps) -> recip * 1/(1+eps)
            CrystalLattice cand = best.Multiply(gemmi::Mat33(s, 0, 0, 0, s, 0, 0, 0, s));
            const double cc = score(cand);
            if (cc > scale_cc) {
                scale_cc = cc;
                scale_best = cand;
            }
        }
        best = scale_best;
        best_cc = scale_cc;
    }

    data.latt = best;
}

// ---------------------------------------------------------------------------
// Term 1 of the factored likelihood (FACTORED_MODEL.md): the per-reflection
// *intensity* (0th-moment) residual. The profile-fit amplitude J should equal the
// scaled reference J_model = G * exp(-B/4d^2) * partiality * pol * I_ref. One scalar
// residual per reflection, weighted by the model-expected (Fisher) sigma_J. This is
// the scaling residual - integration and scaling become one objective, and the empty
// pixels (which make no residual of their own) stop dominating the fit. With geometry
// and R held fixed, J, partiality and sigma_J are constants, so only G and B are free.
// ---------------------------------------------------------------------------
struct IntensityResidual {
    IntensityResidual(double J, double sigma_J, double partiality, double pol,
                      double I_ref, double inv_4d2)
        : J(J), inv_sigma(1.0 / sigma_J), partiality(partiality), pol(pol),
          I_ref(I_ref), inv_4d2(inv_4d2) {}
    template<typename T>
    bool operator()(const T *const G, const T *const B, T *residual) const {
        const T B_term = ceres::exp(-B[0] * T(inv_4d2));
        const T J_model = G[0] * B_term * T(partiality) * T(pol) * T(I_ref);
        residual[0] = (J_model - T(J)) * T(inv_sigma);
        return true;
    }
    double J, inv_sigma, partiality, pol, I_ref, inv_4d2;
};

template<class T>
void PixelRefine::Run(const T *image,
                      BraggPrediction &prediction,
                      PixelRefineData &data) {
    data.solved = false;
    data.reflections.clear();

    // Global orientation + cell-scale sweep before the local LSQ, to recentre the
    // high-resolution shoeboxes onto signal that small misalignments hide.
    if (data.sweep_orientation)
        SweepOrientationCell(image, prediction, data);

    const double lambda = data.geom.GetWavelength_A();
    const double pixel_size = data.geom.GetPixelSize_mm();

    const BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = static_cast<float>(data.ewald_dist_cutoff),
        .max_hkl = 100,
        .centering = data.centering,
        .bandwidth_sigma = static_cast<float>(data.bandwidth)  // relative Δλ/λ sigma
    };

    const int radius = data.shoebox_radius;
    const int bkg_outer_radius = std::max(radius + 1, data.bkg_outer_radius);

    // Per-reflection polarization correction (raw = true * pol), evaluated once at
    // the predicted spot - same handling as BraggIntegrate2D. Identity if disabled.
    const auto pol_factor = experiment.GetPolarizationFactor();
    auto polarization = [&](double x, double y) -> double {
        if (!pol_factor)
            return 1.0;
        return data.geom.CalcAzIntPolarizationCorr(static_cast<float>(x), static_cast<float>(y),
                                                   pol_factor.value());
    };

    // Bandwidth radial-width^2 (in the code's R = sqrt(2)*sigma convention):
    //   R_bw^2 = (b*lambda)^2 / (2 d^4),   b = relative bandwidth (sigma).
    // A fixed per-reflection constant; data.bandwidth == 0 -> monochromatic no-op.
    const double bw = data.bandwidth;
    auto bandwidth_radial_sq = [&](double d) -> double {
        if (bw <= 0.0 || d <= 0.0)
            return 0.0;
        const double bl = bw * lambda;
        return bl * bl / (2.0 * d * d * d * d);
    };

    // Mutable experiment whose geometry is re-synced from the refined data.geom
    // before each prediction, so shoeboxes track the refined geometry/cell.
    DiffractionExperiment exp_iter = experiment;

    // State retained after the loop for the final reflection extraction.
    std::vector<ReflGroup> groups;
    double beam[2] = {0, 0};
    double dist_mm = data.geom.GetDetectorDistance_mm();
    double detector_rot[2] = {0, 0};
    double latt_vec0[3] = {0, 0, 0}; // orientation (Rodrigues)
    double latt_vec1[3] = {0, 0, 0}; // lengths
    double latt_vec2[3] = {0, 0, 0}; // angles (rad)
    double orient_prior[3] = {0, 0, 0}; // pre-refinement orientation (regularization anchor)

    const bool eval_only = (data.max_iterations <= 0);
    const int n_iter = std::max(1, data.max_iterations);
    for (int iter = 0; iter < n_iter; ++iter) {
        // ---- 1. Re-sync prediction geometry from the (refined) model ----------
        exp_iter.BeamX_pxl(data.geom.GetBeamX_pxl())
                .BeamY_pxl(data.geom.GetBeamY_pxl())
                .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
                .PoniRot1_rad(data.geom.GetPoniRot1_rad())
                .PoniRot2_rad(data.geom.GetPoniRot2_rad());

        const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction);

        // ---- 2. Collect per-reflection shoebox pixels -------------------------
        // GetReflections() returns the full pre-sized buffer; only the first
        // nrefl entries are valid for this image (the rest are stale/zeroed).
        groups.clear();
        const auto &predicted = prediction.GetReflections();

        // Spot-core mask over ALL predicted reflections, so each reflection's
        // local background ignores pixels that belong to a neighbouring spot.
        const auto spot_mask = BuildSpotMask(predicted, nrefl, xpixel, ypixel, radius);

        for (int ri = 0; ri < nrefl; ++ri) {
            const auto &refl = predicted[ri];
            const auto hkl = hkl_key_generator(refl);
            if (!reference_data.contains(hkl))
                continue;

            // Local flat background from the ring around the shoebox (raw counts).
            // No azimuthal fallback: if we cannot estimate a clean local background
            // the reflection is dropped, exactly as BraggIntegrate2D marks it
            // unobserved when fewer than a handful of background pixels survive.
            double Ibkg = 0.0;
            if (!EstimateLocalBackground(image, spot_mask, xpixel, ypixel,
                                         refl.predicted_x, refl.predicted_y,
                                         radius, bkg_outer_radius, Ibkg))
                continue;

            ReflGroup g;
            g.h = refl.h;
            g.k = refl.k;
            g.l = refl.l;
            g.d = refl.d;
            g.Itrue = reference_data[hkl];
            g.R_bw_sq = bandwidth_radial_sq(refl.d);
            g.pol = polarization(refl.predicted_x, refl.predicted_y);
            g.Ibkg = Ibkg;
            g.predicted_x = refl.predicted_x;
            g.predicted_y = refl.predicted_y;

            const auto box = ShoeboxBounds(refl.predicted_x, refl.predicted_y, radius, xpixel, ypixel);

            for (int y = box.min_y; y <= box.max_y; ++y) {
                for (int x = box.min_x; x <= box.max_x; ++x) {
                    const size_t npixel = xpixel * y + x;

                    // Skip sentinel (masked / saturated) pixels. We assume the pixel
                    // mask is already applied upstream (encoded as the sentinel).
                    if (image[npixel] == std::numeric_limits<T>::max())
                        continue;
                    if (std::is_signed_v<T> && (image[npixel] == std::numeric_limits<T>::min()))
                        continue;

                    const double Iobs = static_cast<double>(image[npixel]); // raw counts

                    // Variance for the fit weight. Weighting by the observed count
                    // (var = Iobs) lets down-fluctuated background pixels carry the
                    // largest 1/sqrt(var) weight, which biases the fit towards "no
                    // signal" and drove the per-image scale G to 0 on weak images
                    // (collapsing the merge). Use the local background as the
                    // (background-limited) variance, constant over the shoebox - the
                    // same de-biasing applied to the extraction.
                    double var = std::max(Ibkg, 1.0);
                    double weight = 1.0 / std::sqrt(var);

                    // Signal-weighting: down-weight pixels far from the predicted spot
                    // centre so the empty shoebox corners cannot dilute or destabilise
                    // the fit; the signal-bearing core drives the refined parameters.
                    if (data.fit_signal_sigma_pix > 0.0) {
                        const double dx = x - g.predicted_x;
                        const double dy = y - g.predicted_y;
                        const double s2 = data.fit_signal_sigma_pix * data.fit_signal_sigma_pix;
                        weight *= std::exp(-0.5 * (dx * dx + dy * dy) / s2);
                    }

                    PixelObs obs{
                        .x = static_cast<double>(x),
                        .y = static_cast<double>(y),
                        .Iobs = Iobs,
                        .Ibkg = Ibkg,
                        .weight = weight
                    };
                    g.pixels.push_back(obs);
                }
            }

            if (!g.pixels.empty())
                groups.push_back(std::move(g));
        }

        if (groups.empty())
            return;

        // ---- 3. Set up parameter blocks (geometry part mirrors XtalOptimizer) -
        BuildParameterBlocks(data, beam, dist_mm, detector_rot,
                             latt_vec0, latt_vec1, latt_vec2);

        // Anchor for orientation regularization = the orientation the LSQ starts from
        // (captured before the predict<->refine iterations move it). When the global
        // sweep ran first this is the swept orientation, not the original spot-centroid
        // one - which is intended: the regularizer keeps the LSQ near its own starting
        // point, it is not meant to pull a deliberate sweep back.
        if (iter == 0)
            for (int i = 0; i < 3; ++i)
                orient_prior[i] = latt_vec0[i];

        // ---- Term 3: per-reflection recentre on the observed centroid ----------------
        // The geometry predicts the spot to ~0.4 px (per-reflection scatter a global fit
        // cannot remove); a tight Term-2 template centred on the prediction then sits off
        // the real spot. For confident spots, shift the profile centre to the observed
        // centroid (used consistently by Term 2, Term 1 and the extraction below). Weak
        // spots keep the prediction (recentring on a noise centroid would bias them).
        for (auto &g : groups) {
            g.dcx = 0.0;
            g.dcy = 0.0;
        }
        if (data.recenter_profile && !groups.empty()) {
            const double box_px = (2.0 * radius + 1.0) * (2.0 * radius + 1.0);
            for (auto &g : groups) {
                double sw = 0.0, swx = 0.0, swy = 0.0;
                for (const auto &px : g.pixels) {
                    const double w = std::max(px.Iobs - g.Ibkg, 0.0);
                    sw += w; swx += w * px.x; swy += w * px.y;
                }
                if (sw <= 0.0)
                    continue;
                if (sw / std::sqrt(std::max(box_px * g.Ibkg, 1.0)) < data.recenter_min_signif)
                    continue;
                double dcx = swx / sw - g.predicted_x, dcy = swy / sw - g.predicted_y;
                const double dl = std::sqrt(dcx * dcx + dcy * dcy);
                if (dl > 2.0) { dcx *= 2.0 / dl; dcy *= 2.0 / dl; } // clamp runaway centroids
                g.dcx = dcx;
                g.dcy = dcy;
            }
        }

        // ---- Term 2: per-resolution tangential profile width R1 from spot shapes ------
        // Default: every reflection uses the global R1; with shape_R1 on, override it with
        // R1 = sqrt(2*<eps_t^2>) from the intensity-weighted second moment of the strong
        // spots, binned by resolution (low res small spots, high res larger). A shape
        // statistic - normalised by the total, so decoupled from the per-image scale.
        for (auto &g : groups)
            g.R1_eff = data.R[1];
        if (data.shape_R1 && !groups.empty()) {
            constexpr int n_bins = 6;
            const double box_px = (2.0 * radius + 1.0) * (2.0 * radius + 1.0);
            double s2min = 1e30, s2max = 0.0;
            for (const auto &g : groups) {
                const double s2 = 1.0 / (g.d * g.d);
                s2min = std::min(s2min, s2);
                s2max = std::max(s2max, s2);
            }
            const double span = std::max(s2max - s2min, 1e-12);
            auto bin_of = [&](double d) {
                return std::clamp(static_cast<int>((1.0 / (d * d) - s2min) / span * n_bins), 0, n_bins - 1);
            };
            std::vector<std::vector<double>> bin_M2(n_bins);
            for (const auto &g : groups) {
                double sw = 0.0, sw_et2 = 0.0;
                for (const auto &px : g.pixels) {
                    PixelObs probe{px.x - g.dcx, px.y - g.dcy, 0.0, g.Ibkg, 1.0};
                    PixelResidual pr(probe, 1.0, lambda, pixel_size, g.h, g.k, g.l,
                                     g.R_bw_sq, g.pol, data.crystal_system);
                    double q_sq, eps_r, eps_t_sq;
                    if (!pr.GeometryTerms(beam, &dist_mm, detector_rot,
                                          latt_vec0, latt_vec1, latt_vec2, q_sq, eps_r, eps_t_sq))
                        continue;
                    const double w = std::max(px.Iobs - g.Ibkg, 0.0);
                    sw += w;
                    sw_et2 += w * eps_t_sq;
                }
                if (sw <= 0.0)
                    continue;
                const double signif = sw / std::sqrt(std::max(box_px * g.Ibkg, 1.0));
                if (signif >= 5.0) // only well-measured spots define the shape
                    bin_M2[bin_of(g.d)].push_back(sw_et2 / sw);
            }
            std::vector<double> bin_R1(n_bins, data.R[1]);
            for (int b = 0; b < n_bins; ++b)
                if (bin_M2[b].size() >= 5) {
                    const double r1 = std::sqrt(2.0 * std::max(MedianInPlace(bin_M2[b]), 0.0));
                    if (std::isfinite(r1) && r1 > 1e-4)
                        bin_R1[b] = std::clamp(r1, 1e-4, 0.05);
                }
            for (auto &g : groups)
                g.R1_eff = bin_R1[bin_of(g.d)];
            data.shape_R1_lores = bin_R1[0];          // lowest resolution bin
            data.shape_R1_hires = bin_R1[n_bins - 1]; // highest resolution bin
        }

        // ---- 4. Build the problem ---------------------------------------------
        // One residual block per shoebox (N residuals), so the expensive
        // per-reflection node geometry is evaluated once per reflection instead
        // of once per pixel.
        ceres::Problem problem;
        size_t residual_pixels = 0;
        if (data.intensity_residual) {
            // Term-1 path: one per-reflection intensity residual. Geometry & R fixed, so
            // J / partiality / sigma_J are computed here as constants and only G, B vary.
            const double R0 = data.R[0];
            for (const auto &g : groups) {
                const double R1 = g.R1_eff; // Term 2: per-resolution profile width
                double num = 0.0, den = 0.0, rad = 0.0;
                std::vector<std::pair<double, double>> pt_sig; // (P_t, Iobs-Bg) for Fisher pass
                pt_sig.reserve(g.pixels.size());
                for (const auto &px : g.pixels) {
                    PixelObs probe{px.x - g.dcx, px.y - g.dcy, 0.0, g.Ibkg, 1.0}; // Term 3 recentre
                    PixelResidual pr(probe, 1.0, lambda, pixel_size, g.h, g.k, g.l,
                                     g.R_bw_sq, g.pol, data.crystal_system);
                    double q_sq, eps_r, eps_t_sq;
                    if (!pr.GeometryTerms(beam, &dist_mm, detector_rot,
                                          latt_vec0, latt_vec1, latt_vec2, q_sq, eps_r, eps_t_sq))
                        continue;
                    if (!(R1 > 0.0) || !(R0 > 0.0))
                        continue;
                    const double P_t = std::exp(-eps_t_sq / (R1 * R1)) / (M_PI * R1 * R1);
                    const double R0_eff_sq = R0 * R0 + g.R_bw_sq;
                    const double P_rad = std::exp(-eps_r * eps_r / R0_eff_sq);
                    const double v = std::max(g.Ibkg, 1.0);
                    const double sig = px.Iobs - g.Ibkg;
                    num += P_t * sig / v;
                    den += P_t * P_t / v;
                    rad += P_rad * P_t * P_t / v;
                    pt_sig.emplace_back(P_t, sig);
                }
                if (!(den > 0.0))
                    continue;
                const double J = num / den;
                const double partiality = rad / den;
                // Model-expected (Fisher) variance: v_p = background + expected signal J*P_t,
                // not the per-pixel observed counts (which down-bias) - so the weight tracks
                // information, and an expected-strong reflection that is absent hurts.
                double den_f = 0.0;
                for (const auto &[P_t, sig] : pt_sig) {
                    const double v_f = std::max(g.Ibkg + std::max(J, 0.0) * P_t, 1.0);
                    den_f += P_t * P_t / v_f;
                }
                const double sigma_J = std::sqrt(1.0 / std::max(den_f, 1e-30));
                const double inv_4d2 = (g.d > 0.0) ? 1.0 / (4.0 * g.d * g.d) : 0.0;
                auto *cost = new ceres::AutoDiffCostFunction<IntensityResidual, 1, 1, 1>(
                    new IntensityResidual(J, sigma_J, partiality, g.pol, g.Itrue, inv_4d2));
                problem.AddResidualBlock(cost, nullptr, &data.scale_factor, &data.B_factor);
                ++residual_pixels;
            }
            data.residual_count = residual_pixels;
        } else {
        for (const auto &g : groups) {
            auto *cost = new ceres::DynamicAutoDiffCostFunction<ShoeboxResidual>(
                new ShoeboxResidual(g, lambda, pixel_size, data.crystal_system));
            cost->AddParameterBlock(2); // beam
            cost->AddParameterBlock(1); // distance
            cost->AddParameterBlock(2); // detector_rot
            cost->AddParameterBlock(3); // p0 (orientation)
            cost->AddParameterBlock(3); // p1 (lengths)
            cost->AddParameterBlock(3); // p2 (angles)
            cost->AddParameterBlock(1); // scale G
            cost->AddParameterBlock(1); // B
            cost->AddParameterBlock(2); // R
            cost->SetNumResiduals(static_cast<int>(g.pixels.size()));
            // No robust loss here: a per-block (whole-shoebox) Huber would act on
            // the sum of ~N squared residuals and mis-scale, unlike the previous
            // per-pixel Huber. Per-pixel sigma weighting is retained; per-pixel
            // outlier rejection (zingers) is a TODO if needed.
            problem.AddResidualBlock(cost, nullptr,
                                     beam, &dist_mm, detector_rot,
                                     latt_vec0, latt_vec1, latt_vec2,
                                     &data.scale_factor, &data.B_factor, data.R);
            residual_pixels += g.pixels.size();
        }
        data.residual_count = residual_pixels;
        }

        // ---- 5. Constrain / bound parameter blocks ----------------------------
        if (data.intensity_residual) {
            // Only G and B are in this problem; geometry/R are not parameters here.
            problem.SetParameterLowerBound(&data.scale_factor, 0, 0.0);
            if (!data.refine_B)
                problem.SetParameterBlockConstant(&data.B_factor);
            // Regularize G->1, weight sqrt(n_refl/sigma): commensurate because the data
            // term is now one residual per reflection (unlike the per-pixel path).
            if (data.scale_reg_sigma > 0.0 && !groups.empty()) {
                const double w = std::sqrt(static_cast<double>(groups.size()) / data.scale_reg_sigma);
                auto *reg = new ceres::AutoDiffCostFunction<ScalarRegularizer, 1, 1>(
                    new ScalarRegularizer(w, 1.0));
                problem.AddResidualBlock(reg, nullptr, &data.scale_factor);
            }
        } else {
        if (!data.refine_orientation) {
            problem.SetParameterBlockConstant(latt_vec0);
        } else if (data.orient_reg_sigma_deg > 0.0) {
            // Anchor orientation to its spot-centroid prior. The weight is scaled to
            // the *pixel* data term (sqrt(n_pixels)/sigma_rad), not the reflection
            // count - the data has one residual per shoebox pixel, so a reflection-
            // scaled prior (~50x too weak) was simply not felt. At a misorientation of
            // orient_reg_sigma_deg the prior matches the data, so the fit only moves
            // further when the pixels strongly agree it should.
            const double sigma_rad = std::max(data.orient_reg_sigma_deg * M_PI / 180.0, 1e-9);
            const double w = std::sqrt(static_cast<double>(residual_pixels)) / sigma_rad;
            auto *reg = new ceres::AutoDiffCostFunction<OrientationRegularizer, 3, 3>(
                new OrientationRegularizer(w, orient_prior));
            problem.AddResidualBlock(reg, nullptr, latt_vec0);
        }

        if (!data.refine_unit_cell) {
            problem.SetParameterBlockConstant(latt_vec1);
            problem.SetParameterBlockConstant(latt_vec2);
        } else {
            for (int i = 0; i < 3; ++i) {
                problem.SetParameterLowerBound(latt_vec1, i, 5.0);
                problem.SetParameterUpperBound(latt_vec1, i, 1000.0);
            }
            if (data.crystal_system != gemmi::CrystalSystem::Monoclinic &&
                data.crystal_system != gemmi::CrystalSystem::Triclinic)
                problem.SetParameterBlockConstant(latt_vec2);
        }

        if (!data.refine_beam_center)
            problem.SetParameterBlockConstant(beam);

        if (!data.refine_distance) {
            problem.SetParameterBlockConstant(&dist_mm);
        } else {
            problem.SetParameterLowerBound(&dist_mm, 0, dist_mm * 0.9);
            problem.SetParameterUpperBound(&dist_mm, 0, dist_mm * 1.1);
        }

        if (!data.refine_detector_angles) {
            problem.SetParameterBlockConstant(detector_rot);
        } else {
            const double rng = 3.0 / 180.0 * M_PI;
            for (int i = 0; i < 2; ++i) {
                problem.SetParameterLowerBound(detector_rot, i, detector_rot[i] - rng);
                problem.SetParameterUpperBound(detector_rot, i, detector_rot[i] + rng);
            }
        }

        if (data.refine_scale) {
            problem.SetParameterLowerBound(&data.scale_factor, 0, 0.0);
            // Regularize G towards 1 so weakly-constrained images cannot wander
            // (an unconstrained 1/G is what collapsed the cross-image merge). Weight
            // scaled to the pixel data term (n_pixels), not the reflection count.
            if (data.scale_reg_sigma > 0.0) {
                const double w = std::sqrt(static_cast<double>(residual_pixels) / data.scale_reg_sigma);
                auto *reg = new ceres::AutoDiffCostFunction<ScalarRegularizer, 1, 1>(
                    new ScalarRegularizer(w, 1.0));
                problem.AddResidualBlock(reg, nullptr, &data.scale_factor);
            }
        } else {
            problem.SetParameterBlockConstant(&data.scale_factor);
        }

        if (!data.refine_B)
            problem.SetParameterBlockConstant(&data.B_factor);

        if (data.refine_R) {
            if (data.fix_R0) {
                // Diagnostic: hold R0 constant, refine R1 only.
                problem.SetManifold(data.R, new ceres::SubsetManifold(2, {0}));
            } else {
                problem.SetParameterLowerBound(data.R, 0, 1e-5);
                problem.SetParameterLowerBound(data.R, 1, 1e-5);
            }
        } else {
            problem.SetParameterBlockConstant(data.R);
        }
        } // end per-pixel (non-intensity_residual) constraints

        // ---- 6. Solve (or, for max_iterations<=0, just evaluate the cost) -----
        // Evaluate-only is the live-residual path: it reports the current cost
        // without moving any parameter, so a UI can show how good the present
        // R0/R1/bandwidth/geometry are as the user drags sliders.
        if (eval_only) {
            double cost = 0.0;
            problem.Evaluate(ceres::Problem::EvaluateOptions(), &cost, nullptr, nullptr, nullptr);
            data.final_cost = cost;
            data.solved = true;
        } else {
            ceres::Solver::Options options;
            options.linear_solver_type = ceres::DENSE_QR;
            options.minimizer_progress_to_stdout = false;
            options.logging_type = ceres::LoggingType::SILENT;
            options.max_solver_time_in_seconds = data.max_time_s;
            options.num_threads = 1;

            ceres::Solver::Summary summary;
            ceres::Solve(options, &problem, &summary);

            data.final_cost = summary.final_cost;
            data.solved = summary.IsSolutionUsable();

            // Diagnostic: Pearson correlations on the final solve. Always needs G and B
            // refined for G-B; the R correlations are added only when R is also refined.
            if (data.compute_covariance && data.solved && iter == n_iter - 1 &&
                data.refine_scale && data.refine_B) {
                ceres::Covariance::Options copt;
                copt.algorithm_type = ceres::DENSE_SVD;
                copt.null_space_rank = -1; // tolerate (and reveal) degenerate directions
                ceres::Covariance cov(copt);
                std::vector<std::pair<const double *, const double *>> blocks = {
                    {&data.scale_factor, &data.scale_factor}, {&data.B_factor, &data.B_factor},
                    {&data.scale_factor, &data.B_factor}};
                if (data.refine_R) {
                    blocks.push_back({data.R, data.R});
                    blocks.push_back({&data.scale_factor, data.R});
                    blocks.push_back({&data.B_factor, data.R});
                }
                if (cov.Compute(blocks, &problem)) {
                    double cGG, cBB, cGB;
                    cov.GetCovarianceBlock(&data.scale_factor, &data.scale_factor, &cGG);
                    cov.GetCovarianceBlock(&data.B_factor, &data.B_factor, &cBB);
                    cov.GetCovarianceBlock(&data.scale_factor, &data.B_factor, &cGB);
                    const double sG = std::sqrt(std::max(cGG, 0.0));
                    const double sB = std::sqrt(std::max(cBB, 0.0));
                    auto rho = [](double c, double s1, double s2) {
                        return (s1 > 0.0 && s2 > 0.0) ? c / (s1 * s2) : NAN;
                    };
                    data.corr_GB = rho(cGB, sG, sB);
                    if (data.refine_R) {
                        double cRR[4], cGR[2], cBR[2];
                        cov.GetCovarianceBlock(data.R, data.R, cRR);
                        cov.GetCovarianceBlock(&data.scale_factor, data.R, cGR);
                        cov.GetCovarianceBlock(&data.B_factor, data.R, cBR);
                        const double sR0 = std::sqrt(std::max(cRR[0], 0.0));
                        const double sR1 = std::sqrt(std::max(cRR[3], 0.0));
                        data.corr_GR0  = rho(cGR[0], sG, sR0);
                        data.corr_GR1  = rho(cGR[1], sG, sR1);
                        data.corr_BR0  = rho(cBR[0], sB, sR0);
                        data.corr_BR1  = rho(cBR[1], sB, sR1);
                        data.corr_R0R1 = rho(cRR[1], sR0, sR1);
                    }
                    data.covariance_valid = true;
                }
            }
        }

        // ---- 7. Write refined geometry + lattice back into data ---------------
        if (data.refine_beam_center)
            data.geom.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]);
        if (data.refine_distance)
            data.geom.DetectorDistance_mm(dist_mm);
        if (data.refine_detector_angles)
            data.geom.PoniRot1_rad(detector_rot[0]).PoniRot2_rad(detector_rot[1]);

        if (data.refine_orientation || data.refine_unit_cell) {
            switch (data.crystal_system) {
                case gemmi::CrystalSystem::Orthorhombic:
                    data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, M_PI/2);
                    break;
                case gemmi::CrystalSystem::Tetragonal:
                    latt_vec1[1] = latt_vec1[0];
                    data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, M_PI/2);
                    break;
                case gemmi::CrystalSystem::Cubic:
                    latt_vec1[1] = latt_vec1[0];
                    latt_vec1[2] = latt_vec1[0];
                    data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, M_PI/2);
                    break;
                case gemmi::CrystalSystem::Hexagonal:
                    latt_vec1[1] = latt_vec1[0];
                    data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, 2.0*M_PI/3.0);
                    break;
                case gemmi::CrystalSystem::Monoclinic:
                    data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, latt_vec2[0], M_PI/2);
                    break;
                default:
                    data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1,
                                                          latt_vec2[0], latt_vec2[1], latt_vec2[2]);
                    break;
            }
        }
    } // predict<->refine iterations

    // ---- Adaptive integration mask --------------------------------------------
    // Measure R1 (tangential profile width) from the intensity-weighted tangential
    // second moment of the strong spots, rather than fitting it. R1 is a *shape*
    // statistic: sigma_t^2 = sum_p (I_p-B) eps_t,p^2 / sum_p (I_p-B), normalised by the
    // total so it is independent of the per-image scale - which is exactly what breaks
    // the R1<->G degeneracy (a measured width cannot be traded against G). One value per
    // image here (from the strong, mostly low-res spots); a per-resolution version is the
    // natural next step for the high-res / DMM-streak shapes.
    if (data.adaptive_R1 && !groups.empty()) {
        std::vector<double> itrue;
        itrue.reserve(groups.size());
        for (const auto &g : groups)
            itrue.push_back(g.Itrue);
        const size_t cut_idx = itrue.size() * 7 / 10; // keep the strongest ~30%
        std::nth_element(itrue.begin(), itrue.begin() + cut_idx, itrue.end());
        const double itrue_cut = itrue[cut_idx];

        std::vector<double> sigma_t2;
        for (const auto &g : groups) {
            if (g.Itrue < itrue_cut)
                continue;
            double sw = 0.0, sw_et2 = 0.0;
            for (const auto &px : g.pixels) {
                PixelObs probe{px.x, px.y, 0.0, g.Ibkg, 1.0};
                PixelResidual pr(probe, 1.0, lambda, pixel_size, g.h, g.k, g.l,
                                 g.R_bw_sq, g.pol, data.crystal_system);
                double q_sq, eps_r, eps_t_sq;
                if (!pr.GeometryTerms(beam, &dist_mm, detector_rot,
                                      latt_vec0, latt_vec1, latt_vec2, q_sq, eps_r, eps_t_sq))
                    continue;
                const double w = std::max(px.Iobs - g.Ibkg, 0.0);
                sw += w;
                sw_et2 += w * eps_t_sq;
            }
            if (sw > 0.0)
                sigma_t2.push_back(sw_et2 / sw);
        }
        if (sigma_t2.size() >= 5) {
            const double r1 = std::sqrt(2.0 * MedianInPlace(sigma_t2)); // R1^2 = 2 sigma_t^2
            if (std::isfinite(r1) && r1 > 1e-5)
                data.R[1] = r1;
        }
    }

    // ---- Centering diagnostic --------------------------------------------------
    // Observed-centroid vs predicted-position offset for the strong spots, after all
    // refinement. Large rms (relative to the spot size) means a tight profile mask
    // sits off the real spot - which is why a generous box can beat profile fitting.
    if (data.measure_centroid && !groups.empty()) {
        const double beam_x = data.geom.GetBeamX_pxl();
        const double beam_y = data.geom.GetBeamY_pxl();
        const double box_px = (2.0 * radius + 1.0) * (2.0 * radius + 1.0);
        // Raw pixel value (sentinel/bounds safe) for the parabolic peak fit. A constant
        // background cancels in the parabola, so no need to subtract it here.
        auto pix_val = [&](int x, int y) -> double {
            if (x < 0 || x >= static_cast<int>(xpixel) || y < 0 || y >= static_cast<int>(ypixel))
                return std::numeric_limits<double>::quiet_NaN();
            const size_t np = static_cast<size_t>(xpixel) * y + x;
            if (image[np] == std::numeric_limits<T>::max())
                return std::numeric_limits<double>::quiet_NaN();
            if (std::is_signed_v<T> && image[np] == std::numeric_limits<T>::min())
                return std::numeric_limits<double>::quiet_NaN();
            return static_cast<double>(image[np]);
        };
        struct Off { double signif, tang_c, tang_p, rad_c; };
        std::vector<Off> offs;
        double sdx = 0.0, sdy = 0.0, sd2 = 0.0;
        size_t nc = 0;
        for (const auto &g : groups) {
            double sw = 0.0, swx = 0.0, swy = 0.0, bmax = -1e30;
            int bx = 0, by = 0;
            for (const auto &px : g.pixels) {
                const double s = px.Iobs - g.Ibkg;
                const double w = std::max(s, 0.0);
                sw += w; swx += w * px.x; swy += w * px.y;
                if (s > bmax) { bmax = s; bx = static_cast<int>(std::lround(px.x)); by = static_cast<int>(std::lround(px.y)); }
            }
            if (sw <= 0.0)
                continue;
            const double signif = sw / std::sqrt(std::max(box_px * g.Ibkg, 1.0));
            if (signif < 5.0)
                continue; // a measurable spot at any resolution (not just the strong low-res ones)
            // Sub-pixel peak (mode): parabola through the brightest pixel and its two
            // neighbours per axis. The mode tracks the prediction even when an asymmetric
            // tail drags the centroid (mean) sideways - so peak vs centroid separates a
            // shape asymmetry from a true position error.
            double peak_x = bx, peak_y = by;
            { const double l = pix_val(bx - 1, by), c = pix_val(bx, by), r = pix_val(bx + 1, by);
              const double den = l - 2.0 * c + r;
              if (std::isfinite(den) && den < -1e-9)
                  peak_x = bx + std::clamp(0.5 * (l - r) / den, -1.0, 1.0); }
            { const double l = pix_val(bx, by - 1), c = pix_val(bx, by), r = pix_val(bx, by + 1);
              const double den = l - 2.0 * c + r;
              if (std::isfinite(den) && den < -1e-9)
                  peak_y = by + std::clamp(0.5 * (l - r) / den, -1.0, 1.0); }
            const double dcx = swx / sw - g.predicted_x, dcy = swy / sw - g.predicted_y;
            const double dpx = peak_x - g.predicted_x,   dpy = peak_y - g.predicted_y;
            sdx += dcx; sdy += dcy; sd2 += dcx * dcx + dcy * dcy; ++nc;
            const double rx = g.predicted_x - beam_x, ry = g.predicted_y - beam_y;
            const double rr = std::sqrt(rx * rx + ry * ry);
            if (rr < 1.0)
                continue;
            const double rad_c  = (dcx * rx + dcy * ry) / rr;  // signed radial (outward +)
            const double tang_c = (dcx * -ry + dcy * rx) / rr; // signed tangential
            const double tang_p = (dpx * -ry + dpy * rx) / rr;
            offs.push_back({signif, std::fabs(tang_c), std::fabs(tang_p), rad_c});
        }
        if (nc >= 5) {
            data.centroid_bias_px = std::sqrt((sdx / nc) * (sdx / nc) + (sdy / nc) * (sdy / nc));
            data.centroid_rms_px = std::sqrt(sd2 / nc);
        }
        if (offs.size() >= 10) {
            std::vector<double> sig;
            sig.reserve(offs.size());
            for (const auto &o : offs)
                sig.push_back(o.signif);
            std::nth_element(sig.begin(), sig.begin() + sig.size() / 2, sig.end());
            const double smed = sig[sig.size() / 2];
            double slo = 0, shi = 0, tclo = 0, tchi = 0, tplo = 0, tphi = 0, rclo = 0, rchi = 0;
            int nlo = 0, nhi = 0;
            for (const auto &o : offs) {
                if (o.signif < smed) { slo += o.signif; tclo += o.tang_c; tplo += o.tang_p; rclo += o.rad_c; ++nlo; }
                else                 { shi += o.signif; tchi += o.tang_c; tphi += o.tang_p; rchi += o.rad_c; ++nhi; }
            }
            if (nlo > 0 && nhi > 0) {
                data.centroid_lo_signif = slo / nlo;   data.centroid_hi_signif = shi / nhi;
                data.centroid_lo_tang_c = tclo / nlo;  data.centroid_hi_tang_c = tchi / nhi;
                data.centroid_lo_tang_p = tplo / nlo;  data.centroid_hi_tang_p = tphi / nhi;
                data.centroid_lo_rad_c  = rclo / nlo;  data.centroid_hi_rad_c  = rchi / nhi;
            }
        }
    }

    // ---- Extract integrated reflections ---------------------------------------
    // Profile fitting gives the recorded amplitude (fitting the tangential profile
    // P_t against the background-subtracted pixels):
    //     J      = sum_p[ P_t,p (Iobs_p - Ibkg)/v_p ] / sum_p[ P_t,p^2 / v_p ]
    //            ~ G * Itrue * B_term * partiality * pol     (recorded raw counts)
    //     var(J) = 1 / sum_p[ P_t,p^2 / v_p ]
    //
    // Two SEPARATE fractions reduce the full intensity to what these pixels record:
    //
    //   partiality   - the radial / rocking dimension that a still does NOT sample.
    //                  Only the slice of the reflection that crosses the Ewald
    //                  sphere on this shot is recorded; <= 1. We DIVIDE it out to
    //                  recover the full intensity. = profile-weighted P_radial.
    //
    //   completeness - the fraction of the spot's detector footprint that landed on
    //                  live pixels (= profile captured by live pixels / profile over
    //                  the whole shoebox). 1.0 when the spot sits fully on the
    //                  detector; < 1.0 only when a detector edge, gap or mask clips
    //                  it. Profile fitting already extrapolates over the missing
    //                  pixels, so this is NOT applied to r.I - it is a quality flag.
    //
    // Output split (Merge multiplies r.I * image_scale_corr and weights by
    // 1/(sigma*image_scale_corr)^2 - see Merge.cpp):
    //   r.I              = J / (B_term * partiality * pol) = G * Itrue
    //   r.sigma          = sqrt(var(J)) / (B_term * partiality * pol)
    //   r.partiality     = profile-weighted P_radial in (0,1]  (the rocking fraction)
    //   r.completeness   = live/total tangential profile in (0,1]  (detector clipping)
    //   r.image_scale_corr = 1/G                         (per-image scale ONLY)
    // so r.I * image_scale_corr = Itrue. B, partiality and polarization live on the
    // intensity, G lives on image_scale_corr - one clean meaning per field.
    //
    // We walk the full (unclamped) shoebox once: every grid point feeds the total
    // tangential profile (completeness denominator); points that are real, live
    // detector pixels also feed the profile fit and the captured profile.
    data.reflections.reserve(groups.size());
    for (const auto &g : groups) {
        const int cx = static_cast<int>(std::lround(g.predicted_x));
        const int cy = static_cast<int>(std::lround(g.predicted_y));

        // Debye-Waller factor for this reflection (constant over its shoebox).
        const double B_term = std::exp(-data.B_factor / (4.0 * g.d * g.d));

        // Term 3 recentre shift (precomputed in the pre-pass; 0 if off or the spot was not
        // confident). Profile evaluated at (x-dcx, y-dcy), data summed at (x,y).
        const double dcx = g.dcx, dcy = g.dcy;

        double num = 0.0, den = 0.0, bkg_sum = 0.0, radial_sum = 0.0;
        double prof_live = 0.0, prof_full = 0.0; // tangential profile: captured / total
        size_t n = 0;

        for (int y = cy - radius; y <= cy + radius; ++y) {
            for (int x = cx - radius; x <= cx + radius; ++x) {
                // Geometry/profile for this grid point (profile recentred by (dcx,dcy)).
                PixelObs probe{static_cast<double>(x) - dcx, static_cast<double>(y) - dcy,
                               0.0, g.Ibkg, 1.0};
                PixelResidual pr(probe, 1.0, lambda, pixel_size, g.h, g.k, g.l,
                                 g.R_bw_sq, g.pol, data.crystal_system);
                double q_sq, eps_r, eps_t_sq;
                if (!pr.GeometryTerms(beam, &dist_mm, detector_rot,
                                      latt_vec0, latt_vec1, latt_vec2, q_sq, eps_r, eps_t_sq))
                    continue;
                if (!(data.R[0] > 0.0) || !(g.R1_eff > 0.0))
                    continue;

                // Tangential profile shape (area-normalized) -> the fit template. Uses the
                // per-reflection R1_eff (Term 2), falling back to the global R1 by default.
                const double R1 = g.R1_eff;
                const double P_t = std::exp(-eps_t_sq / (R1 * R1)) / (M_PI * R1 * R1);
                prof_full += P_t; // whole shoebox, on- or off-detector

                // Only real, unmasked detector pixels carry signal.
                if (x < 0 || x >= static_cast<int>(xpixel) || y < 0 || y >= static_cast<int>(ypixel))
                    continue;
                const size_t np = static_cast<size_t>(xpixel) * y + x;
                if (image[np] == std::numeric_limits<T>::max())
                    continue;
                if (std::is_signed_v<T> && image[np] == std::numeric_limits<T>::min())
                    continue;

                const double Iobs = static_cast<double>(image[np]); // raw counts
                // Variance for the profile-fit weights. Weighting by the *observed*
                // per-pixel count (v = Iobs) biases the amplitude negative: a
                // down-fluctuated background pixel gets the smallest v and hence the
                // largest 1/v weight, so num = sum P_t (Iobs - Ibkg)/v is pulled below
                // zero - worst where the signal is weakest, i.e. the high-resolution
                // shells (the negative <I/sig> we see there). For background-limited
                // reflections the variance is the local background, constant over the
                // shoebox, so use that instead of the observed count.
                double v = std::max(g.Ibkg, 1.0);

                // Peak-normalized radial factor (the partiality), in (0,1]. The
                // bandwidth-broadened radial width matches the model in Model().
                const double R0_eff_sq = data.R[0] * data.R[0] + g.R_bw_sq;
                const double P_radial = std::exp(-eps_r * eps_r / R0_eff_sq);

                // Profile-fit accumulators. The amplitude estimator weights pixels by
                // P_t^2/v, so the partiality (which de-scales that amplitude) MUST use
                // the SAME weights - otherwise an R0_eff-dependent (resolution-
                // dependent) factor is left behind in r.I.
                const double w = P_t * P_t / v;
                num += P_t * (Iobs - g.Ibkg) / v;
                den += w;
                radial_sum += P_radial * w;   // partiality weighted exactly like num/den
                prof_live += P_t;             // captured tangential profile
                bkg_sum += g.Ibkg;
                ++n;
            }
        }

        Reflection r{};
        r.h = g.h;
        r.k = g.k;
        r.l = g.l;
        r.d = static_cast<float>(g.d);
        r.predicted_x = static_cast<float>(g.predicted_x);
        r.predicted_y = static_cast<float>(g.predicted_y);
        r.observed_x = NAN;
        r.observed_y = NAN;
        r.rlp = 1.0f;
        r.partiality = (den > 0.0) ? static_cast<float>(radial_sum / den) : 1.0f;
        r.completeness = (prof_full > 0.0) ? static_cast<float>(prof_live / prof_full) : 1.0f;

        if (den > 0.0 && n > 0) {
            const double I_amp = num / den;             // ~ G*Itrue*B_term*partiality*pol
            const double sigma_amp = std::sqrt(1.0 / den);
            const double corr = static_cast<double>(r.partiality) * B_term * g.pol; // B, partiality & pol
            r.bkg = static_cast<float>(bkg_sum / static_cast<double>(n));
            r.observed = true;

            if (corr > 0.0 && data.scale_factor > 0.0) {
                r.I = static_cast<float>(I_amp / corr);       // = G*Itrue
                r.sigma = static_cast<float>(sigma_amp / corr);
                r.image_scale_corr = static_cast<float>(1.0 / data.scale_factor); // G only
            } else {
                r.I = static_cast<float>(I_amp);
                r.sigma = static_cast<float>(sigma_amp);
                r.image_scale_corr = NAN;
            }
        } else {
            r.I = 0.0f;
            r.sigma = NAN;
            r.bkg = 0.0f;
            r.observed = false;
        }
        data.reflections.push_back(r);
    }

    // ---- Per-image CC vs reference (the half/ref correlation diagnostic) -------
    // Pearson CC of the scaled estimate (r.I * image_scale_corr = Itrue_est)
    // against the reference intensities, over the matched reflections.
    {
        double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
        size_t cn = 0;
        for (const auto &r : data.reflections) {
            if (!r.observed || !std::isfinite(r.I) || !std::isfinite(r.image_scale_corr))
                continue;
            const auto it = reference_data.find(hkl_key_generator(r));
            if (it == reference_data.end())
                continue;
            const double x = static_cast<double>(r.I) * static_cast<double>(r.image_scale_corr);
            const double y = it->second;
            if (!std::isfinite(x) || !std::isfinite(y))
                continue;
            sx += x; sy += y; sxx += x * x; syy += y * y; sxy += x * y; ++cn;
        }
        data.cc = NAN;
        data.cc_n = static_cast<int64_t>(cn);
        if (cn >= 3) {
            const double nd = static_cast<double>(cn);
            const double cov = sxy - sx * sy / nd;
            const double vx = sxx - sx * sx / nd;
            const double vy = syy - sy * sy / nd;
            if (vx > 0.0 && vy > 0.0)
                data.cc = cov / std::sqrt(vx * vy);
        }
    }
}

template<class T>
std::vector<float> PixelRefine::PredictImage(const T *image,
                                             BraggPrediction &prediction,
                                             const PixelRefineData &data,
                                             bool include_background) const {
    std::vector<float> img(xpixel * ypixel, 0.0f);

    const double lambda = data.geom.GetWavelength_A();
    const double pixel_size = data.geom.GetPixelSize_mm();
    const int radius = data.shoebox_radius;
    const int bkg_outer_radius = std::max(radius + 1, data.bkg_outer_radius);
    const double bw = data.bandwidth;

    const auto pol_factor = experiment.GetPolarizationFactor();
    auto polarization = [&](double x, double y) -> double {
        if (!pol_factor)
            return 1.0;
        return data.geom.CalcAzIntPolarizationCorr(static_cast<float>(x), static_cast<float>(y),
                                                   pol_factor.value());
    };
    auto bandwidth_radial_sq = [&](double d) -> double {
        if (bw <= 0.0 || d <= 0.0)
            return 0.0;
        const double bl = bw * lambda;
        return bl * bl / (2.0 * d * d * d * d);
    };

    double beam[2], dist_mm, detector_rot[2];
    double latt_vec0[3], latt_vec1[3], latt_vec2[3];
    BuildParameterBlocks(data, beam, dist_mm, detector_rot,  latt_vec0, latt_vec1, latt_vec2);

    DiffractionExperiment exp_iter = experiment;
    exp_iter.BeamX_pxl(data.geom.GetBeamX_pxl())
            .BeamY_pxl(data.geom.GetBeamY_pxl())
            .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
            .PoniRot1_rad(data.geom.GetPoniRot1_rad())
            .PoniRot2_rad(data.geom.GetPoniRot2_rad());

    const BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = static_cast<float>(data.ewald_dist_cutoff),
        .max_hkl = 100,
        .centering = data.centering,
        .bandwidth_sigma = static_cast<float>(data.bandwidth)  // relative Δλ/λ sigma
    };
    const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction);
    const auto &predicted = prediction.GetReflections();
    const auto spot_mask = BuildSpotMask(predicted, nrefl, xpixel, ypixel, radius);

    for (int ri = 0; ri < nrefl; ++ri) {
        const auto &refl = predicted[ri];
        const auto it = reference_data.find(hkl_key_generator(refl));
        if (it == reference_data.end())
            continue;

        const double Itrue = it->second;
        const double R_bw_sq = bandwidth_radial_sq(refl.d);
        const double pol = polarization(refl.predicted_x, refl.predicted_y);

        // Local background straight from the actual image (flat per shoebox), laid
        // into the box so the prediction overlays the real frame - the same model
        // path Run() fits, now reproduced faithfully because we have the image.
        double Ibkg = 0.0;
        const bool have_bkg = include_background &&
            EstimateLocalBackground(image, spot_mask, xpixel, ypixel,
                                    refl.predicted_x, refl.predicted_y,
                                    radius, bkg_outer_radius, Ibkg);

        const auto box = ShoeboxBounds(refl.predicted_x, refl.predicted_y, radius, xpixel, ypixel);

        for (int y = box.min_y; y <= box.max_y; ++y) {
            for (int x = box.min_x; x <= box.max_x; ++x) {
                const size_t npixel = xpixel * y + x;

                PixelObs obs{
                    .x = static_cast<double>(x),
                    .y = static_cast<double>(y),
                    .Iobs = 0.0,
                    .Ibkg = have_bkg ? Ibkg : 0.0,
                    .weight = 1.0
                };
                PixelResidual pr(obs, Itrue, lambda, pixel_size,
                                 refl.h, refl.k, refl.l, R_bw_sq, pol, data.crystal_system);

                double Ipred = 0.0; // raw counts: signal (+ local background)
                if (pr.Model(beam, &dist_mm, detector_rot,
                             latt_vec0, latt_vec1, latt_vec2,
                             &data.scale_factor, &data.B_factor, data.R, Ipred))
                    img[npixel] += static_cast<float>(Ipred);
            }
        }
    }

    return img;
}

template<class T>
std::vector<float> PixelRefine::ChiSquaredImage(const T *image,
                                                BraggPrediction &prediction,
                                                const PixelRefineData &data) const {
    std::vector<float> img(xpixel * ypixel, 0.0f);

    const double lambda = data.geom.GetWavelength_A();
    const double pixel_size = data.geom.GetPixelSize_mm();
    const int radius = data.shoebox_radius;
    const int bkg_outer_radius = std::max(radius + 1, data.bkg_outer_radius);
    const double bw = data.bandwidth;

    const auto pol_factor = experiment.GetPolarizationFactor();
    auto polarization = [&](double x, double y) -> double {
        if (!pol_factor)
            return 1.0;
        return data.geom.CalcAzIntPolarizationCorr(static_cast<float>(x), static_cast<float>(y),
                                                   pol_factor.value());
    };
    auto bandwidth_radial_sq = [&](double d) -> double {
        if (bw <= 0.0 || d <= 0.0)
            return 0.0;
        const double bl = bw * lambda;
        return bl * bl / (2.0 * d * d * d * d);
    };

    double beam[2], dist_mm, detector_rot[2];
    double latt_vec0[3], latt_vec1[3], latt_vec2[3];
    BuildParameterBlocks(data, beam, dist_mm, detector_rot,  latt_vec0, latt_vec1, latt_vec2);

    DiffractionExperiment exp_iter = experiment;
    exp_iter.BeamX_pxl(data.geom.GetBeamX_pxl())
            .BeamY_pxl(data.geom.GetBeamY_pxl())
            .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
            .PoniRot1_rad(data.geom.GetPoniRot1_rad())
            .PoniRot2_rad(data.geom.GetPoniRot2_rad());

    const BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = static_cast<float>(data.ewald_dist_cutoff),
        .max_hkl = 100,
        .centering = data.centering,
        .bandwidth_sigma = static_cast<float>(data.bandwidth)
    };
    const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction);
    const auto &predicted = prediction.GetReflections();
    const auto spot_mask = BuildSpotMask(predicted, nrefl, xpixel, ypixel, radius);

    for (int ri = 0; ri < nrefl; ++ri) {
        const auto &refl = predicted[ri];
        const auto it = reference_data.find(hkl_key_generator(refl));
        if (it == reference_data.end())
            continue;

        const double Itrue = it->second;
        const double R_bw_sq = bandwidth_radial_sq(refl.d);
        const double pol = polarization(refl.predicted_x, refl.predicted_y);

        // Local flat background, identical to Run(); skip the reflection if it
        // cannot be estimated (matches Run() dropping the reflection).
        double Ibkg = 0.0;
        if (!EstimateLocalBackground(image, spot_mask, xpixel, ypixel,
                                     refl.predicted_x, refl.predicted_y,
                                     radius, bkg_outer_radius, Ibkg))
            continue;

        const auto box = ShoeboxBounds(refl.predicted_x, refl.predicted_y, radius, xpixel, ypixel);

        for (int y = box.min_y; y <= box.max_y; ++y) {
            for (int x = box.min_x; x <= box.max_x; ++x) {
                const size_t npixel = xpixel * y + x;

                // Same gating as Run(): only pixels that actually enter the fit.
                if (image[npixel] == std::numeric_limits<T>::max())
                    continue;
                if (std::is_signed_v<T> && (image[npixel] == std::numeric_limits<T>::min()))
                    continue;

                const double Iobs = static_cast<double>(image[npixel]); // raw counts

                double var = std::max(Iobs, 0.0);
                if (!(var > 1.0))
                    var = 1.0;
                const double weight = 1.0 / std::sqrt(var);

                PixelObs obs{
                    .x = static_cast<double>(x),
                    .y = static_cast<double>(y),
                    .Iobs = Iobs,
                    .Ibkg = Ibkg,
                    .weight = weight
                };
                PixelResidual pr(obs, Itrue, lambda, pixel_size,
                                 refl.h, refl.k, refl.l, R_bw_sq, pol, data.crystal_system);

                double Ipred = 0.0;
                if (pr.Model(beam, &dist_mm, detector_rot,
                             latt_vec0, latt_vec1, latt_vec2,
                             &data.scale_factor, &data.B_factor, data.R, Ipred)) {
                    // residual_i = (I_pred - I_obs) * weight (== Ceres residual);
                    // its square is this pixel's contribution to the cost.
                    const double rw = (Ipred - Iobs) * weight;
                    img[npixel] += static_cast<float>(rw * rw);
                }
            }
        }
    }

    return img;
}

// Explicit instantiations for the supported (uncompressed) image pixel types.
template void PixelRefine::Run<int8_t>(const int8_t *, BraggPrediction &, PixelRefineData &);
template void PixelRefine::Run<int16_t>(const int16_t *, BraggPrediction &, PixelRefineData &);
template void PixelRefine::Run<int32_t>(const int32_t *, BraggPrediction &, PixelRefineData &);
template void PixelRefine::Run<uint8_t>(const uint8_t *, BraggPrediction &, PixelRefineData &);
template void PixelRefine::Run<uint16_t>(const uint16_t *, BraggPrediction &, PixelRefineData &);
template void PixelRefine::Run<uint32_t>(const uint32_t *, BraggPrediction &, PixelRefineData &);

template std::vector<float> PixelRefine::PredictImage<int8_t>(const int8_t *, BraggPrediction &, const PixelRefineData &, bool) const;
template std::vector<float> PixelRefine::PredictImage<int16_t>(const int16_t *, BraggPrediction &, const PixelRefineData &, bool) const;
template std::vector<float> PixelRefine::PredictImage<int32_t>(const int32_t *, BraggPrediction &, const PixelRefineData &, bool) const;
template std::vector<float> PixelRefine::PredictImage<uint8_t>(const uint8_t *, BraggPrediction &, const PixelRefineData &, bool) const;
template std::vector<float> PixelRefine::PredictImage<uint16_t>(const uint16_t *, BraggPrediction &, const PixelRefineData &, bool) const;
template std::vector<float> PixelRefine::PredictImage<uint32_t>(const uint32_t *, BraggPrediction &, const PixelRefineData &, bool) const;

template std::vector<float> PixelRefine::ChiSquaredImage<int8_t>(const int8_t *, BraggPrediction &, const PixelRefineData &) const;
template std::vector<float> PixelRefine::ChiSquaredImage<int16_t>(const int16_t *, BraggPrediction &, const PixelRefineData &) const;
template std::vector<float> PixelRefine::ChiSquaredImage<int32_t>(const int32_t *, BraggPrediction &, const PixelRefineData &) const;
template std::vector<float> PixelRefine::ChiSquaredImage<uint8_t>(const uint8_t *, BraggPrediction &, const PixelRefineData &) const;
template std::vector<float> PixelRefine::ChiSquaredImage<uint16_t>(const uint16_t *, BraggPrediction &, const PixelRefineData &) const;
template std::vector<float> PixelRefine::ChiSquaredImage<uint32_t>(const uint32_t *, BraggPrediction &, const PixelRefineData &) const;