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

#include "PixelRefine.h"

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

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

namespace {

// Per-pixel observation, in *corrected* intensity units (solid-angle and
// polarization correction already folded in, consistently for signal and
// background). Geometry-independent quantities are precomputed here so that the
// Ceres cost functor stays cheap.
struct PixelObs {
    double x, y;          // detector pixel coordinate
    double Iobs;          // corrected pixel value (signal + background)
    double Ibkg;          // corrected background estimate (azimuthal bin mean)
    double weight;        // 1 / sigma_pixel
    double A_recip;       // reciprocal-space area subtended by the pixel (Jacobian)
    double angle_rad;     // goniometer angle of this observation
};

// 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 predicted_x, predicted_y;
    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;
}

} // namespace

// ---------------------------------------------------------------------------
// Cost functor
//
//   I_pred(pixel) = G * Itrue * B_term * P_radial * P_tangential + 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 = A_recip/(pi R1^2) * exp(-eps_t^2/R1^2)(spatial profile on the
//                                                        detector, normalized so
//                                                        that sum over pixels ~ 1)
//
// The tangential factor is what makes this "profile fitting": summing
// I_pred - I_bkg over the shoebox reproduces G * Itrue * B_term * P_radial.
// The 1/(pi R1^2) normalization is the missing piece that decouples the profile
// width R1 from 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,
                  gemmi::CrystalSystem symmetry)
        : Itrue(Itrue), Iobs(obs.Iobs), Ibkg(obs.Ibkg), weight(obs.weight),
          A_recip(obs.A_recip), 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), angle_rad(obs.angle_rad), 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 rotation_axis,
                       const T *const p0,
                       const T *const p1,
                       const T *const p2,
                       T &q_sq, T &eps_radial, T &eps_tang_sq) const {
        // PyFAI convention (left-handed for rot1/rot2):
        // poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1), rot3 = 0 assumed.
        const T rot1 = detector_rot[0];
        const T rot2 = detector_rot[1];

        const T c1 = ceres::cos(rot1);
        const T s1 = ceres::sin(rot1);
        const T c2 = ceres::cos(rot2);
        const T s2 = ceres::sin(rot2);

        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 lab_norm = ceres::sqrt(x * x + y * y + z * z);
        const T inv_norm = T(1) / lab_norm;

        T recip_raw[3];
        recip_raw[0] = x * inv_norm * T(inv_lambda);
        recip_raw[1] = y * inv_norm * T(inv_lambda);
        recip_raw[2] = (z * inv_norm - T(1.0)) * T(inv_lambda);

        // Goniometer "back-to-start" rotation: image frame -> reference frame.
        const T aa_back[3] = {
            T(angle_rad) * rotation_axis[0],
            T(angle_rad) * rotation_axis[1],
            T(angle_rad) * rotation_axis[2]
        };
        T recip_obs[3];
        ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs);
        const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);

        // Build cell lengths and the (unit) B matrix from the symmetry-specific
        // parametrization (identical convention to XtalOptimizer::XtalResidual).
        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 V = A.dot(BxC);
        if (ceres::abs(V) < T(1e-12))
            return false;
        const T invV = T(1) / V;

        const Eigen::Matrix<T, 3, 1> Astar = BxC * invV;
        const Eigen::Matrix<T, 3, 1> Bstar = CxA * invV;
        const Eigen::Matrix<T, 3, 1> Cstar = AxB * invV;

        const Eigen::Matrix<T, 3, 1> e_pred_recip =
            Astar * T(exp_h) + Bstar * T(exp_k) + Cstar * T(exp_l);

        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;

        const Eigen::Matrix<T, 3, 1> n_radial = S_pred / S_pred_norm;
        const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;

        eps_radial = delta_q.dot(n_radial);
        const Eigen::Matrix<T, 3, 1> dq_tang = delta_q - eps_radial * n_radial;
        eps_tang_sq = dq_tang.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 rotation_axis,
               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, rotation_axis,
                           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));

        // Full 3D reciprocal-space spot density modelled as a separable Gaussian,
        // normalized so that its integral is 1 (intensity-conserving):
        //   radial:     g_r(e) = exp(-e^2/R0^2) / (sqrt(pi) R0)     [1/A^-1]
        //   tangential: g_t(e) = exp(-|e|^2/R1^2) / (pi R1^2)        [1/A^-2]
        // The detector pixel captures the fraction g_t * A_recip of the tangential
        // profile (A_recip = reciprocal area the pixel subtends; sum over shoebox
        // ~ 1). The radial factor is the still-image partiality: how far this
        // reflection sits from the Ewald sphere.
        //
        // Caveat: a still samples the radial direction at a single offset, so the
        // sqrt(pi) R0 normalization makes g_r a density (1/A^-1) rather than a
        // dimensionless fraction. The leftover dimensional factor is absorbed by
        // the free scale G. The energy-bandwidth contribution to the radial width
        // is folded in here via R_bw_sq (beam divergence is still TODO).
        const T R0_eff_sq = R[0] * R[0] + T(R_bw_sq);
        const T g_radial = ceres::exp(-eps_radial * eps_radial / R0_eff_sq)
                           / (ceres::sqrt(T(M_PI)) * ceres::sqrt(R0_eff_sq));
        const T P_tang = T(A_recip) * 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 * g_radial * P_tang;
        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 rotation_axis,
                    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, rotation_axis,
                   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, A_recip;
    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 angle_rad;
    gemmi::CrystalSystem symmetry;
};

PixelRefine::PixelRefine(const DiffractionExperiment &experiment,
                         const AzimuthalIntegrationMapping &mapping,
                         const std::vector<MergedReflection> &reference)
    : mapping(mapping),
      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;
}

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

    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(),
        .max_hkl = 100,
        .centering = data.centering
    };

    const auto azim_result = profile.GetResult();
    const auto azim_std = profile.GetStd();
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const auto &corrections = mapping.Corrections();
    const int azim_bin_count = mapping.GetAzimuthalBinCount();

    const double angle_rad = data.angle_deg * M_PI / 180.0;
    const int radius = data.shoebox_radius;

    // Exact reciprocal-space area a 1x1 pixel subtends, |dq/dx x dq/dy|, via
    // finite differences of the detector->reciprocal map. This is the Jacobian
    // between the curved Ewald-sphere sampling and flat reciprocal space, and it
    // is exactly the geometric factor that plays the role of the Lorentz factor
    // for stills: where the sphere grazes reciprocal space obliquely, a pixel
    // covers more reciprocal volume and the captured fraction grows. It tracks
    // the refined geometry because it reads the current data.geom each iteration.
    auto recip_area = [&](double x, double y) -> double {
        const Coord qx = data.geom.DetectorToRecip(x + 0.5, y) - data.geom.DetectorToRecip(x - 0.5, y);
        const Coord qy = data.geom.DetectorToRecip(x, y + 0.5) - data.geom.DetectorToRecip(x, y - 0.5);
        return (qx % qy).Length();
    };

    // 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 rot_vec[3] = {1.0, 0.0, 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)

    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());

        prediction.Calc(exp_iter, data.latt, settings_prediction);

        // ---- 2. Collect per-reflection shoebox pixels -------------------------
        groups.clear();
        for (const auto &refl : prediction.GetReflections()) {
            const auto hkl = hkl_key_generator(refl);
            if (!reference_data.contains(hkl))
                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.predicted_x = refl.predicted_x;
            g.predicted_y = refl.predicted_y;

            const int min_y = std::max<int>(refl.predicted_y - radius, 0);
            const int max_y = std::min<int>(refl.predicted_y + radius, ypixel - 1);
            const int min_x = std::max<int>(refl.predicted_x - radius, 0);
            const int max_x = std::min<int>(refl.predicted_x + radius, xpixel - 1);

            for (int y = min_y; y <= max_y; ++y) {
                for (int x = min_x; x <= max_x; ++x) {
                    const size_t npixel = xpixel * y + x;
                    const int azim_bin = pixel_to_bin[npixel];

                    // Skip pixels not mapped to an azimuthal bin or carrying a
                    // sentinel (masked / saturated) value. We assume the pixel
                    // mask is already applied upstream.
                    if (azim_bin >= azim_bin_count)
                        continue;
                    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 correction = corrections[npixel];
                    const double Ibkg = azim_result[azim_bin]; // already in corrected units
                    const double Ibkg_sigma = azim_std[azim_bin];
                    const double raw = static_cast<double>(image[npixel]);
                    const double Iobs = raw * correction;

                    // Per-pixel variance: Poisson noise of the corrected counts
                    // (var(c*N) = c^2 * N = c * Iobs) plus the background spread.
                    double var = correction * std::max(Iobs, 0.0) + Ibkg_sigma * Ibkg_sigma;
                    if (!(var > 1.0))
                        var = 1.0;

                    PixelObs obs{
                        .x = static_cast<double>(x),
                        .y = static_cast<double>(y),
                        .Iobs = Iobs,
                        .Ibkg = Ibkg,
                        .weight = 1.0 / std::sqrt(var),
                        .A_recip = recip_area(x, y),
                        .angle_rad = angle_rad
                    };
                    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) -
        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();
        rot_vec[0] = 1.0; rot_vec[1] = 0.0; rot_vec[2] = 0.0;
        if (auto axis = data.geom.GetRotation()) {
            rot_vec[0] = axis->GetAxis().x;
            rot_vec[1] = axis->GetAxis().y;
            rot_vec[2] = axis->GetAxis().z;
        }

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

        // ---- 4. Build the problem ---------------------------------------------
        ceres::Problem problem;
        for (const auto &g : groups) {
            for (const auto &obs : g.pixels) {
                auto *cost = new ceres::AutoDiffCostFunction<
                    PixelResidual, 1, 2, 1, 2, 3, 3, 3, 3, 1, 1, 2>(
                    new PixelResidual(obs, g.Itrue, lambda, pixel_size,
                                      g.h, g.k, g.l, g.R_bw_sq, data.crystal_system));
                problem.AddResidualBlock(cost, new ceres::HuberLoss(3.0),
                                         beam, &dist_mm, detector_rot, rot_vec,
                                         latt_vec0, latt_vec1, latt_vec2,
                                         &data.scale_factor, &data.B_factor, data.R);
            }
        }
        data.residual_count = problem.NumResidualBlocks();

        // ---- 5. Constrain / bound parameter blocks ----------------------------
        if (!data.refine_orientation)
            problem.SetParameterBlockConstant(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_rotation_axis)
            problem.SetParameterBlockConstant(rot_vec);

        if (data.refine_scale)
            problem.SetParameterLowerBound(&data.scale_factor, 0, 0.0);
        else
            problem.SetParameterBlockConstant(&data.scale_factor);

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

        if (data.refine_R) {
            problem.SetParameterLowerBound(data.R, 0, 1e-5);
            problem.SetParameterLowerBound(data.R, 1, 1e-5);
        } else {
            problem.SetParameterBlockConstant(data.R);
        }

        // ---- 6. Solve ---------------------------------------------------------
        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();

        // ---- 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

    // ---- Extract integrated reflections ---------------------------------------
    // Two quantities are read back per reflection, using the *optimized* model:
    //
    //  * Recorded (partial) intensity J by profile fitting against the normalized
    //    tangential profile P_t (sum over the shoebox ~ 1):
    //        J       = sum_p [ P_t,p (Iobs_p - Ibkg_p) / v_p ] / sum_p [ P_t,p^2 / v_p ]
    //        var(J)  = 1 / sum_p [ P_t,p^2 / v_p ]
    //    Because P_t is area-normalized, J estimates the integrated intensity
    //    actually recorded on this image (not the full-spot intensity).
    //
    //  * Partiality: the profile-weighted mean of the *peak-normalized* radial
    //    factor exp(-eps_r^2/R0^2) in (0,1], i.e. how close to the Ewald sphere
    //    the reflection sits. Kept dimensionless for consistency with the rest of
    //    the pipeline (image_scale_corr = rlp / partiality, full I = J / partiality).
    data.reflections.reserve(groups.size());
    for (const auto &g : groups) {
        double num = 0.0, den = 0.0, bkg_sum = 0.0;
        double radial_sum = 0.0, radial_w = 0.0;
        size_t n = 0;

        for (const auto &obs : g.pixels) {
            PixelResidual pr(obs, 1.0, lambda, pixel_size, g.h, g.k, g.l, g.R_bw_sq, data.crystal_system);
            double q_sq, eps_r, eps_t_sq;
            if (!pr.GeometryTerms(beam, &dist_mm, detector_rot, rot_vec,
                                  latt_vec0, latt_vec1, latt_vec2, q_sq, eps_r, eps_t_sq))
                continue;
            if (!(data.R[0] > 0.0) || !(data.R[1] > 0.0))
                continue;

            // Normalized tangential profile (sum over shoebox ~ 1) -> fit weight.
            const double P_t = obs.A_recip * std::exp(-eps_t_sq / (data.R[1] * data.R[1]))
                               / (M_PI * data.R[1] * data.R[1]);
            // Peak-normalized radial factor (the partiality), in (0,1].
            // Bandwidth-broadened radial width, matching 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);

            const double v = SafeInv(obs.weight * obs.weight, 1.0); // pixel variance
            const double signal = obs.Iobs - obs.Ibkg;

            num += P_t * signal / v;
            den += P_t * P_t / v;
            radial_sum += P_radial * P_t;   // weight partiality by the spot core
            radial_w += P_t;
            bkg_sum += obs.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 = (radial_w > 0.0) ? static_cast<float>(radial_sum / radial_w) : 1.0f;

        if (den > 0.0 && n > 0) {
            r.I = static_cast<float>(num / den);
            r.sigma = static_cast<float>(std::sqrt(1.0 / den));
            r.bkg = static_cast<float>(bkg_sum / static_cast<double>(n));
            r.observed = true;
            // Put I onto a common (merge-ready) scale: I * image_scale_corr is the
            // full-spot, scale/DW-corrected intensity. Fold in the refined G, the
            // per-reflection B_term and the partiality (rlp = 1 here).
            const double B_term = std::exp(-data.B_factor / (4.0 * g.d * g.d));
            const double denom = static_cast<double>(r.partiality) * data.scale_factor * B_term;
            r.image_scale_corr = (denom > 0.0) ? static_cast<float>(r.rlp / denom) : NAN;
        } else {
            r.I = 0.0f;
            r.sigma = NAN;
            r.bkg = 0.0f;
            r.observed = false;
        }
        data.reflections.push_back(r);
    }
}

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