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

#include "PostRefine.h"

#include <algorithm>
#include <array>
#include <cmath>

#include "../../common/JFJochMath.h"   // PI
#include "XtalResidual.h"              // XtalResidual (the positional detector<->reciprocal residual, step B)
#include "LatticeReduction.h"
#include "ceres/ceres.h"
#include "ceres/rotation.h"

namespace {

// One integrated partial, flattened across all images.
struct Partial {
    int h, k, l;
    float img;
    double I, sigma;
    double angle_rad;   // frame mid-exposure goniometer angle
    double obs_x, obs_y; // observed spot centroid (pixels); NAN if the box sum found no centroid
};

// A rocking event and its precomputed reference reciprocal vector (phi=0 frame, from the indexed lattice).
struct Event {
    double phi_obs;     // rad, intensity-weighted rocking centroid
    double weight;      // sqrt(sum I / sum sigma)
    double e_ref[3];    // h*a* + k*b* + l*c* at the reference (unrefined) cell/orientation
    int h, k, l;
};

// Distance-INDEPENDENT Ewald excitation residual for a uniform cell-scale parameter s and a refined
// goniometer axis (3-vector). The header-distance miscalibration leaves a uniform cell scale; the axis is
// the other phi_obs lever. Both are phi_obs-constrained (distance-independent). e_ref is the reference
// reciprocal (h*a* + k*b* + l*c* at the indexed cell). On the Ewald sphere <=> |p|^2 + 2 p_z/lambda == 0.
struct ScaleAxisExcitationResidual {
    ScaleAxisExcitationResidual(double lambda, double angle_rad, double weight, const double e_ref[3])
        : inv_lambda(1.0 / lambda), angle_rad(angle_rad), weight(weight),
          ex(e_ref[0]), ey(e_ref[1]), ez(e_ref[2]) {}
    template<typename T>
    bool operator()(const T *const s, const T *const axis, T *residual) const {
        const T inv_s = T(1) / s[0];
        const T p_ref[3] = {T(ex) * inv_s, T(ey) * inv_s, T(ez) * inv_s};
        const T aa[3] = {T(-angle_rad) * axis[0], T(-angle_rad) * axis[1], T(-angle_rad) * axis[2]};
        T p_lab[3];
        ceres::AngleAxisRotatePoint(aa, p_ref, p_lab);
        const T zeta = p_lab[0] * p_lab[0] + p_lab[1] * p_lab[1] + p_lab[2] * p_lab[2]
                       + T(2.0) * p_lab[2] * T(inv_lambda);
        residual[0] = T(weight) * zeta * T(0.5) / T(inv_lambda);
        return true;
    }
    const double inv_lambda, angle_rad, weight, ex, ey, ez;
};

// GONIOMETER ROTATION SCALE k: the same Ewald excitation residual, but with the cell scale and the axis
// DIRECTION already committed by step A, so the single free quantity is how far the stage actually turned
// per unit of commanded angle. Two differences from step A matter:
//   * the angle is measured from the CENTRE of the sweep, not from the goniometer's zero. The reference
//     orientation is the one rotation indexing fitted against the commanded angles, so it has already
//     absorbed the MEAN angle error; only the part that varies across the sweep is left to fit. Scaling the
//     absolute angle instead - which is what reading k off the length of step A's axis vector does - asks
//     the fit to also produce a constant offset it has no parameter for, and the least-squares compromise
//     shrinks k towards 1 by var(phi) / (var(phi) + phi_centre^2): exactly a factor of four for the common
//     case of a sweep starting at zero.
//   * e_mid is the reference reciprocal vector already turned to the sweep centre and divided by the
//     committed cell scale, so nothing but k is free.
struct RotationScaleResidual {
    RotationScaleResidual(double lambda, double dangle_rad, const double u[3], const double e_mid[3])
        : inv_lambda(1.0 / lambda), dangle_rad(dangle_rad),
          ux(u[0]), uy(u[1]), uz(u[2]), ex(e_mid[0]), ey(e_mid[1]), ez(e_mid[2]) {}
    template<typename T>
    bool operator()(const T *const k, T *residual) const {
        const T a = T(-dangle_rad) * k[0];
        const T aa[3] = {a * T(ux), a * T(uy), a * T(uz)};
        const T p_ref[3] = {T(ex), T(ey), T(ez)};
        T p_lab[3];
        ceres::AngleAxisRotatePoint(aa, p_ref, p_lab);
        const T zeta = p_lab[0] * p_lab[0] + p_lab[1] * p_lab[1] + p_lab[2] * p_lab[2]
                       + T(2.0) * p_lab[2] * T(inv_lambda);
        residual[0] = zeta * T(0.5) / T(inv_lambda);
        return true;
    }
    const double inv_lambda, dangle_rad, ux, uy, uz, ex, ey, ez;
};

}  // namespace

PostRefineResult PostRefineRotationGeometry(const std::vector<IntegrationOutcome> &outcomes,
                                            const GoniometerAxis &axis,
                                            const DiffractionGeometry &nominal_geom,
                                            const CrystalLattice &reference_latt,
                                            const PostRefineSettings &settings,
                                            Logger &logger) {
    PostRefineResult result;
    result.geom = nominal_geom;
    result.cell = reference_latt.GetUnitCell();
    result.distance_before_mm = nominal_geom.GetDetectorDistance_mm();
    result.distance_after_mm = nominal_geom.GetDetectorDistance_mm();
    try {
        const double wedge_half = axis.GetWedge_deg() / 2.0;
        const double lambda = nominal_geom.GetWavelength_A();
        const Coord ax = axis.GetAxis();
        const Coord Astar = reference_latt.Astar(), Bstar = reference_latt.Bstar(), Cstar = reference_latt.Cstar();

        std::vector<Partial> pts;
        for (const auto &o : outcomes)
            for (const auto &r : o.reflections) {
                if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f) continue;
                const double mid_deg = axis.GetAngle_deg(r.image_number) + wedge_half;
                const double ox = std::isfinite(r.observed_x) ? r.observed_x : NAN;
                const double oy = std::isfinite(r.observed_y) ? r.observed_y : NAN;
                pts.push_back(Partial{r.h, r.k, r.l, r.image_number, r.I, r.sigma,
                                      mid_deg * PI / 180.0, ox, oy});
            }
        logger.Info("Post-refine: {} partials gathered", pts.size());
        if (pts.size() < static_cast<size_t>(settings.min_events)) return result;

        std::sort(pts.begin(), pts.end(), [](const Partial &a, const Partial &b) {
            if (a.h != b.h) return a.h < b.h;
            if (a.k != b.k) return a.k < b.k;
            if (a.l != b.l) return a.l < b.l;
            return a.img < b.img;
        });

        // Split into rocking events (same raw hkl, adjacent frames). Only >=2-frame events carry an
        // unbiased phi_obs (a single-frame centroid is just the frame centre); precompute e_ref per event.
        constexpr float MAX_FRAME_GAP = 2.0f;
        std::vector<Event> events;
        size_t i = 0, event_frames = 0;
        while (i < pts.size()) {
            size_t j = i + 1;
            while (j < pts.size() && pts[j].h == pts[i].h && pts[j].k == pts[i].k && pts[j].l == pts[i].l
                   && pts[j].img - pts[j - 1].img <= MAX_FRAME_GAP)
                ++j;
            if (j - i >= 2) {
                double sumI = 0, sumIphi = 0, sumSig = 0;
                for (size_t m = i; m < j; ++m) {
                    const double Ipos = std::max(0.0, pts[m].I);
                    sumI += Ipos; sumIphi += Ipos * pts[m].angle_rad; sumSig += pts[m].sigma;
                }
                if (sumI > 0.0 && sumSig > 0.0) {
                    const double phi = sumIphi / sumI;
                    event_frames += j - i;
                    const Coord e = Astar * static_cast<float>(pts[i].h) + Bstar * static_cast<float>(pts[i].k)
                                    + Cstar * static_cast<float>(pts[i].l);
                    events.push_back(Event{phi, std::sqrt(sumI / sumSig), {e.x, e.y, e.z},
                                           pts[i].h, pts[i].k, pts[i].l});
                }
            }
            i = j;
        }
        // Frames per event is the phi_obs sampling: near 2 the reflections barely rock, so the angle
        // this refinement is fitted to is under-determined. It is a geometry count, so unlike an
        // intensity-weighted width it cannot be inflated by noise.
        logger.Info("Post-refine: {} multi-frame rocking events ({:.1f} frames per event)", events.size(),
                    events.empty() ? 0.0 : static_cast<double>(event_frames) / events.size());
        if (static_cast<int>(events.size()) < settings.min_events) return result;


        // The rotation-scale fit further down is a single scalar whose whole point is how the residual
        // varies ALONG the sweep, so it keeps every event. The cap below ranks by I/sigma, and on the
        // crystals that have a stage fault the strong events sit in the middle of the sweep - the part
        // that still indexes - so a capped set would leave the ends unrepresented in exactly the fit that
        // has to see them.
        const std::vector<Event> scale_events = events;

        constexpr size_t MAX_EVENTS = 20000;
        if (events.size() > MAX_EVENTS) {
            std::nth_element(events.begin(), events.begin() + MAX_EVENTS, events.end(),
                             [](const Event &a, const Event &b) { return a.weight > b.weight; });
            events.resize(MAX_EVENTS);
        }

        // ---- GEOMETRY REFINEMENT: the XtalOptimizer-equivalent, done as TWO SEPARATE
        // cross-validated steps rather than one joint fit (the same lesson as integration: refining the
        // profile width and the scale jointly fails, refining them separately works). Each step is committed
        // only if it lowers a HELD-OUT (deterministic split-half) residual - otherwise that part of the
        // geometry is left at nominal ("quit when things go wrong"):
        //   Step A: cell scale + rotation axis from phi_obs (distance-independent excitation residual).
        //   Step B: detector distance + beam centre from the observed spot positions, with the cell FIXED at
        //           step A (so the positional residual is no longer degenerate with the cell scale).
        // Detector tilt is held fixed (gauge-coupled to orientation on a single crystal). ----
        if (settings.refine_geometry) {
            const gemmi::CrystalSystem sys =
                (settings.crystal_system == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal
                                                                            : settings.crystal_system;
            const double ax0[3] = {ax.x, ax.y, ax.z};
            const double lambda_l = lambda;
            const double rot3 = nominal_geom.GetPoniRot3_rad();
            const double pixel_mm = nominal_geom.GetPixelSize_mm();
            const double det_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()};
            const UnitCell r0 = reference_latt.GetUnitCell();

            // Deterministic split of the reflections into a fit half and a held-out half. Avalanche-mix the
            // hkl hash so the split bit is decorrelated from the LSB - a plain h+k+l parity collides with the
            // lattice centering condition (e.g. an I-centred lattice has h+k+l even for EVERY present
            // reflection, so a parity split would leave the validation half empty).
            auto is_val = [](int h, int k, int l) {
                unsigned u = static_cast<unsigned>(h) * 2654435761u + static_cast<unsigned>(k) * 2246822519u
                             + static_cast<unsigned>(l) * 3266489917u;
                u ^= u >> 15; u *= 2246822519u; u ^= u >> 13;
                return (u & 1u) != 0u;
            };
            enum Subset { FIT, VAL, ALL };
            auto in = [&](int h, int k, int l, Subset s) {
                return s == ALL || (is_val(h, k, l) == (s == VAL)); };

            // ===== Step A: cell scale s + rotation axis from phi_obs =====
            auto excit_cost = [&](Subset s, double sc, const double axv[3]) {
                double c = 0.0; int n = 0;
                for (const auto &ev : events) {
                    if (!in(ev.h, ev.k, ev.l, s)) continue;
                    ScaleAxisExcitationResidual r(lambda_l, ev.phi_obs, 1.0, ev.e_ref);
                    double sd = sc, av[3] = {axv[0], axv[1], axv[2]}, resid = 0.0;
                    r(&sd, av, &resid); c += resid * resid; ++n;
                }
                return n ? c / n : 0.0;
            };
            auto solve_scale_axis = [&](Subset s, double &s_out, double ax_out[3]) {
                double sc = 1.0, axv[3] = {ax0[0], ax0[1], ax0[2]};
                ceres::Problem p;
                for (const auto &ev : events) {
                    if (!in(ev.h, ev.k, ev.l, s)) continue;
                    p.AddResidualBlock(new ceres::AutoDiffCostFunction<ScaleAxisExcitationResidual, 1, 1, 3>(
                        new ScaleAxisExcitationResidual(lambda_l, ev.phi_obs, settings.excitation_weight, ev.e_ref)),
                        new ceres::CauchyLoss(0.02), &sc, axv);
                }
                p.SetParameterLowerBound(&sc, 0, 0.9); p.SetParameterUpperBound(&sc, 0, 1.1);
                for (int j = 0; j < 3; ++j) { p.SetParameterLowerBound(axv, j, ax0[j] - 0.05);
                                              p.SetParameterUpperBound(axv, j, ax0[j] + 0.05); }
                ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 50;
                o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
                ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
                s_out = sc; ax_out[0] = axv[0]; ax_out[1] = axv[1]; ax_out[2] = axv[2];
                return sum.IsSolutionUsable();
            };
            double s_fit = 1.0, ax_fit[3];
            const bool convA = solve_scale_axis(FIT, s_fit, ax_fit);
            const double cvA_nom = excit_cost(VAL, 1.0, ax0);
            const double cvA_ref = excit_cost(VAL, s_fit, ax_fit);
            // Commit the cell scale only for a small, credible move: a well-calibrated header needs < ~0.6 %,
            // so a > 1 % scale is a red flag (on multi-lattice / noisy data the excitation fit is biased the
            // same way in every cross-validation fold, so the relative-improvement gate cannot catch it).
            result.cell_refined = convA && cvA_ref < 0.98 * cvA_nom && std::fabs(s_fit - 1.0) < 0.01;
            double s = 1.0, axv[3] = {ax0[0], ax0[1], ax0[2]};
            if (result.cell_refined) solve_scale_axis(ALL, s, axv);   // commit: re-fit on all data
            const double axlen = std::sqrt(axv[0]*axv[0] + axv[1]*axv[1] + axv[2]*axv[2]);
            const double axdev = std::acos(std::clamp((axv[0]*ax0[0]+axv[1]*ax0[1]+axv[2]*ax0[2])
                                                      / std::max(1e-9, axlen), -1.0, 1.0)) * 180.0 / PI;
            logger.Info("Post-refine GEOM step A (cell/axis): s = {:.5f}, rot-axis {:.3f} deg, held-out excit "
                        "{:.3e} -> {:.3e} => {}", s, axdev, cvA_nom, cvA_ref,
                        result.cell_refined ? "COMMIT" : "reject (kept nominal cell)");

            // ===== Goniometer rotation SCALE k, its own one-parameter fit on the same rocking events =====
            // The angles stored in the file are the COMMANDED ones, so a stage that turned k times as far
            // is invisible in the header. Nothing else here can represent it: the cell scale, the axis
            // direction, the distance and the beam are all orthogonal to a rotation MAGNITUDE error. Fitted
            // after step A so the cell scale and the axis direction are fixed at their committed values and
            // k is the only free quantity.
            const double u[3] = {axv[0] / axlen, axv[1] / axlen, axv[2] / axlen};
            double phi_c = 0.0, phi_lo = scale_events[0].phi_obs, phi_hi = scale_events[0].phi_obs;
            for (const auto &ev : scale_events) {
                phi_c += ev.phi_obs;
                phi_lo = std::min(phi_lo, ev.phi_obs);
                phi_hi = std::max(phi_hi, ev.phi_obs);
            }
            phi_c /= static_cast<double>(scale_events.size());
            const double sweep_deg = (phi_hi - phi_lo) * 180.0 / PI;
            // The reference reciprocal vector turned to the sweep centre, at the committed cell scale. The
            // angle then enters the fit measured FROM that centre. A constant crystal missetting about the
            // spindle is k with a slope in phi, so measuring the angle from the goniometer's zero instead
            // lets a missetting leak into k with gain <phi>/<phi^2> - which depends only on where the sweep
            // happens to sit. On a short sweep starting near zero that gain is enormous: a 0.14 deg
            // missetting on a 10 deg wedge fakes 1.4 % of k. Referred to the sweep centre the leak is
            // identically zero at any width, and no parameter has to be added to get it.
            std::vector<std::array<double, 3>> e_mid(scale_events.size());
            for (size_t e = 0; e < scale_events.size(); ++e) {
                const double aa[3] = {-phi_c * u[0], -phi_c * u[1], -phi_c * u[2]};
                const double p[3] = {scale_events[e].e_ref[0] / s, scale_events[e].e_ref[1] / s,
                                     scale_events[e].e_ref[2] / s};
                ceres::AngleAxisRotatePoint(aa, p, e_mid[e].data());
            }
            // Robust-loss scale from the scatter the events actually have: it varies by more than a decade
            // between datasets, so any fixed constant is either inert or throws away real data.
            double rms = 0.0;
            for (size_t e = 0; e < scale_events.size(); ++e) {
                RotationScaleResidual r(lambda_l, scale_events[e].phi_obs - phi_c, u, e_mid[e].data());
                double one = 1.0, resid = 0.0;
                r(&one, &resid); rms += resid * resid;
            }
            rms = std::sqrt(rms / static_cast<double>(scale_events.size()));
            // Fit k over the events whose phi lies outside the given fifth of the sweep (-1 = all of it).
            auto solve_scale = [&](int drop_fifth) {
                double kv = 1.0;
                ceres::Problem p;
                for (size_t e = 0; e < scale_events.size(); ++e) {
                    const int fifth = std::clamp(static_cast<int>(
                            5.0 * (scale_events[e].phi_obs - phi_lo) / std::max(1e-9, phi_hi - phi_lo)), 0, 4);
                    if (fifth == drop_fifth) continue;
                    p.AddResidualBlock(new ceres::AutoDiffCostFunction<RotationScaleResidual, 1, 1>(
                        new RotationScaleResidual(lambda_l, scale_events[e].phi_obs - phi_c, u, e_mid[e].data())),
                        new ceres::HuberLoss(std::max(1e-12, 2.0 * rms)), &kv);
                }
                p.SetParameterLowerBound(&kv, 0, 0.95); p.SetParameterUpperBound(&kv, 0, 1.05);
                ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 50;
                o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
                ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
                return sum.IsSolutionUsable() ? kv : 1.0;
            };
            const double k_fit = solve_scale(-1);
            result.rotation_scale = k_fit;

            // ----- Whether to COMMIT it. A stage fault is rare - 36 of 37 rotation datasets sit at 1.0000
            // on a direct scan - and a 1 % angle correction applied to a healthy dataset would damage it
            // silently, so every test below has to pass.
            // Preconditions: below these the fit is reported but never acted on. Under ~30 deg of sweep k
            // entangles with the axis direction and 10-20 deg truncations of a perfect dataset wander by
            // +-0.6 %; a screening wedge must not trigger a correction.
            constexpr int MIN_SCALE_EVENTS = 5000;
            constexpr double MIN_SCALE_SWEEP_DEG = 30.0;
            // T1 significance: 0.5 % is 18 sigma on the between-dataset scatter of healthy stages
            // (robust sd 2.8e-4) and still 3.5x below the one measured fault.
            constexpr double ROTATION_SCALE_TOL = 0.005;
            // T2 relevance: the misorientation the error produces at each end of the sweep. A large k over
            // a short sweep moves nothing and is not worth correcting.
            constexpr double MIN_SCALE_END_ERROR_DEG = 0.5;
            // T3 uniformity: a stage error is a ramp present in EVERY part of the sweep, so dropping any
            // fifth of it must leave the same k. A second lattice that dominates ONE END of the sweep -
            // exactly what happens where the primary stops indexing - fakes a k indistinguishable from a
            // real fault on T1 and T2, and is the reason this test is not optional. It replaces the
            // hkl-hash split used elsewhere here, which cannot see it: both halves of that split sit at
            // the same angles, so anything structured in phi survives in both folds.
            constexpr double MIN_SCALE_JACKKNIFE_FRAC = 0.5;
            const double end_error_deg = std::fabs(k_fit - 1.0) * sweep_deg / 2.0;
            const bool enough_data = static_cast<int>(scale_events.size()) >= MIN_SCALE_EVENTS
                                     && sweep_deg >= MIN_SCALE_SWEEP_DEG;
            const bool big_enough = enough_data && std::fabs(k_fit - 1.0) >= ROTATION_SCALE_TOL
                                    && end_error_deg >= MIN_SCALE_END_ERROR_DEG;
            double jackknife = 1.0;
            if (big_enough)
                for (int f = 0; f < 5; ++f)
                    jackknife = std::min(jackknife, (solve_scale(f) - 1.0) / (k_fit - 1.0));
            result.rotation_scale_suspect = big_enough && jackknife >= MIN_SCALE_JACKKNIFE_FRAC;
            logger.Info("Post-refine rotation SCALE: k = {:.5f} over {:.0f} deg of sweep centred on {:.1f} "
                        "deg ({} events): end error {:.2f} deg, leave-a-fifth-out {:.2f} => {}",
                        k_fit, sweep_deg, phi_c * 180.0 / PI, scale_events.size(), end_error_deg, jackknife,
                        result.rotation_scale_suspect ? "COMMIT"
                        : !enough_data ? "report only (too little sweep or too few events)"
                        : "reject (kept the stored angles)");
            if (result.rotation_scale_suspect)
                logger.Warning("Goniometer rotation scale looks off by {:+.2f} % (fitted {:.5f}): the stage "
                               "appears to have turned {} than the angles stored in the file, which are the "
                               "COMMANDED values. This is a hardware calibration fault, not a data problem - "
                               "left uncorrected it inflates mosaicity, biases the cell and loses "
                               "high-resolution reflections",
                               100.0 * (k_fit - 1.0), k_fit, k_fit > 1.0 ? "further" : "less far");

            // Cell (scale s, shape fixed) as the XtalResidual parameter blocks p0/p1/p2, held CONSTANT in step B.
            double p0[3] = {0, 0, 0}, p1[3] = {0, 0, 0}, p2[3] = {0, 0, 0};
            double beta = r0.beta;
            switch (sys) {
                case gemmi::CrystalSystem::Tetragonal:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
                    p1[0] = (p1[0] + p1[1]) / 2.0; break;
                case gemmi::CrystalSystem::Cubic:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
                    p1[0] = (p1[0] + p1[1] + p1[2]) / 3.0; break;
                case gemmi::CrystalSystem::Hexagonal:
                    LatticeToRodriguesAndLengths_Hex(reference_latt, p0, p1); break;
                case gemmi::CrystalSystem::Monoclinic:
                    LatticeToRodriguesLengthsBeta_Mono(reference_latt, p0, p1, beta);
                    p2[0] = beta; break;
                case gemmi::CrystalSystem::Orthorhombic:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1); break;
                default:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
                    p2[0] = r0.alpha * PI / 180.0; p2[1] = r0.beta * PI / 180.0; p2[2] = r0.gamma * PI / 180.0; break;
            }
            for (int j = 0; j < 3; ++j) p1[j] *= s;                     // apply the committed cell scale
            double rot_vec[3] = {axv[0], axv[1], axv[2]};              // committed (or nominal) axis

            // ===== Step B: detector distance + beam from the observed positions, cell fixed =====
            std::vector<const Partial *> obs;
            for (const auto &pp : pts)
                if (std::isfinite(pp.obs_x) && std::isfinite(pp.obs_y)) obs.push_back(&pp);
            constexpr size_t MAX_OBS = 20000;
            if (obs.size() > MAX_OBS) {
                std::nth_element(obs.begin(), obs.begin() + MAX_OBS, obs.end(),
                                 [](const Partial *a, const Partial *b) {
                                     return a->I / std::max(1e-9, a->sigma) > b->I / std::max(1e-9, b->sigma); });
                obs.resize(MAX_OBS);
            }
            result.obs_used = static_cast<int>(obs.size());
            const double beam_x0 = nominal_geom.GetBeamX_pxl(), beam_y0 = nominal_geom.GetBeamY_pxl();
            const double dist0 = nominal_geom.GetDetectorDistance_mm();
            auto pos_cost = [&](Subset s, const double beam[2], const double dist[1]) {
                double c = 0.0; int n = 0;
                for (const Partial *pp : obs) {
                    if (!in(pp->h, pp->k, pp->l, s)) continue;
                    XtalResidual r(pp->obs_x, pp->obs_y, lambda_l, pixel_mm, rot3, pp->angle_rad,
                                   pp->h, pp->k, pp->l, sys);
                    double resid[3] = {0, 0, 0};
                    r(beam, dist, det_rot, rot_vec, p0, p1, p2, resid);
                    c += resid[0]*resid[0] + resid[1]*resid[1] + resid[2]*resid[2]; ++n;
                }
                return n ? c / n : 0.0;
            };
            auto solve_detector = [&](Subset s, double beam_out[2], double &dist_out) {
                double beam[2] = {beam_x0, beam_y0}, dist[1] = {dist0};
                ceres::Problem p;
                for (const Partial *pp : obs) {
                    if (!in(pp->h, pp->k, pp->l, s)) continue;
                    p.AddResidualBlock(new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
                        new XtalResidual(pp->obs_x, pp->obs_y, lambda_l, pixel_mm, rot3, pp->angle_rad,
                                         pp->h, pp->k, pp->l, sys)),
                        new ceres::CauchyLoss(0.02), beam, dist,
                        const_cast<double *>(det_rot), rot_vec, p0, p1, p2);
                    p.SetParameterBlockConstant(const_cast<double *>(det_rot));
                    p.SetParameterBlockConstant(rot_vec);
                    p.SetParameterBlockConstant(p0); p.SetParameterBlockConstant(p1); p.SetParameterBlockConstant(p2);
                }
                if (p.NumResidualBlocks() == 0) { beam_out[0] = beam_x0; beam_out[1] = beam_y0; dist_out = dist0; return false; }
                p.SetParameterLowerBound(dist, 0, dist0 * 0.95); p.SetParameterUpperBound(dist, 0, dist0 * 1.05);
                for (int j = 0; j < 2; ++j) { p.SetParameterLowerBound(beam, j, beam[j] - 15.0);
                                              p.SetParameterUpperBound(beam, j, beam[j] + 15.0); }
                ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 60;
                o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
                ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
                beam_out[0] = beam[0]; beam_out[1] = beam[1]; dist_out = dist[0];
                return sum.IsSolutionUsable();
            };
            double beam[2] = {beam_x0, beam_y0}, dist = dist0;
            if (obs.size() >= static_cast<size_t>(settings.min_events)) {
                double beam_fit[2], dist_fit;
                const bool convB = solve_detector(FIT, beam_fit, dist_fit);
                const double b_nom[2] = {beam_x0, beam_y0}, d_nom[1] = {dist0};
                const double b_ref[2] = {beam_fit[0], beam_fit[1]}, d_ref[1] = {dist_fit};
                const double cvB_nom = pos_cost(VAL, b_nom, d_nom);
                const double cvB_ref = pos_cost(VAL, b_ref, d_ref);
                // Commit the detector geometry only for a small, credible move: distance < 1 % (a calibrated
                // header needs < ~0.6 %). A larger move is the red flag for an unreliable fit - typically a
                // second lattice whose spots bias every cross-validation fold identically, so the relative
                // "it improved" gate is blind to it and pulls a spurious distance<->cell pair (the radial
                // degeneracy) far off. The absolute size of the move discriminates a genuine header correction
                // from that failure far better than the absolute residual, which real marginal (noisy / iced)
                // data shares with the multi-lattice case.
                const bool in_bounds = std::fabs(dist_fit - dist0) < 0.01 * dist0
                                       && std::hypot(beam_fit[0] - beam_x0, beam_fit[1] - beam_y0) < 15.0;
                result.detector_refined = convB && cvB_ref < 0.98 * cvB_nom && in_bounds;
                if (result.detector_refined) { double bo[2]; solve_detector(ALL, bo, dist); beam[0] = bo[0]; beam[1] = bo[1]; }
                logger.Info("Post-refine GEOM step B (distance/beam): dist {:.3f} -> {:.3f} mm, beam "
                            "({:.2f},{:.2f}) -> ({:.2f},{:.2f}), held-out pos {:.3e} -> {:.3e} => {}",
                            dist0, result.detector_refined ? dist : dist0, beam_x0, beam_y0,
                            result.detector_refined ? beam[0] : beam_x0, result.detector_refined ? beam[1] : beam_y0,
                            cvB_nom, cvB_ref, result.detector_refined ? "COMMIT" : "reject (kept nominal detector)");
            } else {
                logger.Info("Post-refine GEOM step B: only {} positional observations - skipped", obs.size());
            }

            // Assemble the committed geometry.
            UnitCell cellA = r0;
            if (result.cell_refined) { cellA.a = static_cast<float>(r0.a * s); cellA.b = static_cast<float>(r0.b * s);
                                       cellA.c = static_cast<float>(r0.c * s); }
            result.cell = cellA;
            result.distance_after_mm = dist;
            result.beam_x_before_px = beam_x0; result.beam_x_after_px = beam[0];
            result.beam_y_before_px = beam_y0; result.beam_y_after_px = beam[1];
            result.events_used = static_cast<int>(events.size());
            result.ok = result.cell_refined || result.detector_refined;
            if (!result.ok)
                logger.Info("Post-refine GEOM: neither step passed cross-validation - geometry left at nominal");
            return result;
        }

        return result;   // refine_geometry is the only supported mode; nothing refined otherwise
    } catch (...) {
        result.ok = false;
        return result;
    }
}