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

#include "../../common/JFJochMath.h"
#include <algorithm>
#include <Eigen/Dense>

#include "XtalOptimizer.h"
#include "XtalResidual.h"
#include "XtalRefine.h"
#include "ceres/rotation.h"
#include "Dual.h"
#include "LatticeReduction.h"

// Prior confidence weight per spot: how strong the spot is FOR ITS RESOLUTION. The frame's spots are
// ordered by resolution and cut into equal-count shells, and each intensity is divided by its shell
// median. Refinement needs the high-resolution spots (they carry the cell and distance information) and
// those are legitimately weaker, so a raw intensity weight would suppress exactly the wrong ones; the
// shell normalisation makes the weight resolution-neutral by construction.
//
// The weight enters as w^2 on the squared residual, w^2 = r/(1+r): the shell median contributes half,
// a 4x-median spot 0.8, a quarter-median spot 0.2. Weak spots still pull, they just do not drive. Unlike
// a robust loss this is a PRIOR - it never looks at the current residual, so it cannot mistake a genuine
// spot for an outlier when the starting geometry is far off and leave the fit unable to move.
static std::vector<double> SpotConfidenceWeights(const std::vector<SpotToSave> &spots) {
    constexpr size_t spots_per_shell = 32;

    // Resolution order. Sorting a packed (resolution, index) array rather than an index vector with a
    // projection into the spots keeps the comparisons off the 80-byte records - the same keys in the
    // same order, so introsort makes the same comparisons and the same swaps, and the order it leaves
    // is the same.
    struct SpotByRes {
        float d_A;
        uint32_t index;
    };
    std::vector<SpotByRes> by_res(spots.size());
    for (size_t i = 0; i < spots.size(); i++)
        by_res[i] = {spots[i].d_A, static_cast<uint32_t>(i)};
    std::ranges::sort(by_res, {}, &SpotByRes::d_A);

    const size_t nshells = std::max<size_t>(1, spots.size() / spots_per_shell);
    std::vector<double> weight(spots.size());
    std::vector<float> shell_intensity;

    for (size_t s = 0; s < nshells; s++) {
        const size_t begin = s * spots.size() / nshells;
        const size_t end = (s + 1) * spots.size() / nshells;

        shell_intensity.clear();
        for (size_t i = begin; i < end; i++)
            shell_intensity.push_back(spots[by_res[i].index].intensity);
        std::ranges::nth_element(shell_intensity, shell_intensity.begin() + shell_intensity.size() / 2);
        const double median = std::max(1e-3f, shell_intensity[shell_intensity.size() / 2]);

        for (size_t i = begin; i < end; i++) {
            const double r = std::max(0.0f, spots[by_res[i].index].intensity) / median;
            weight[by_res[i].index] = std::sqrt(r / (1.0 + r));
        }
    }
    return weight;
}

// The oscillation width at which the acceptance gate starts profiling out the rotation coordinate.
// The BAND is principled: the dead zone below matters once the exposure's own rms rotation ambiguity,
// wedge/sqrt(12) = 0.29*wedge, is comparable to the crystal's intrinsic along-u rocking spread, which
// measures ~0.26 deg, and that puts the boundary somewhere between 0.25 and 1.0 deg. The POINT is
// empirical and is taken at the conservative end of that band, because fine slicing is the core case
// and coarse slicing is compatibility: below this the gate is left exactly as it was.
constexpr float COARSE_SLICING_WEDGE_DEG = 0.5f;

// The dead zone's half-width as a fraction of the exposure: the rms of a rotation coordinate uniform
// over the frame, which is the width a least-squares is calibrated on. Half the exposure - the worst
// case a spot could sit at - and forgiving the direction outright were both measured worse.
const double DEAD_ZONE_K = 1.0 / std::sqrt(12.0);

bool XtalOptimizerInternal(XtalOptimizerData &data,
                           std::span<const std::vector<SpotToSave>> spots,
                           const std::vector<std::vector<double>> &weights,
                           const float tolerance,
                           const int num_threads) {
    try {
        // A coplanar basis has no reciprocal cell: 1/V is infinite, every predicted reciprocal vector
        // comes out NaN, and the solver fails on the very first evaluation. There is nothing for the
        // refinement to recover here, so refuse the lattice before the problem is built rather than let
        // the solver discover it. The check has to be on the vectors: this close to flat, float cell
        // angles no longer carry even the SIGN of the metric determinant, and the triclinic branch of
        // XtalResidual then clamps c into the a-b plane and divides by the zero volume that makes.
        if (data.latt.VolumeFraction() < MIN_BASIS_VOLUME_FRACTION)
            return false;

        Coord vec0 = data.latt.Vec0();
        Coord vec1 = data.latt.Vec1();
        Coord vec2 = data.latt.Vec2();
        double beta = data.latt.GetUnitCell().beta;

        // Initial guess for the parameters
        const double distance_mm = data.geom.GetDetectorDistance_mm();

        XtalRefineProblem problem;
        problem.crystal_system = data.crystal_system;
        problem.distance_mm = distance_mm;
        double *beam = problem.beam;
        beam[0] = data.geom.GetBeamX_pxl();
        beam[1] = data.geom.GetBeamY_pxl();
        double *detector_rot = problem.detector_rot;
        detector_rot[0] = data.geom.GetPoniRot1_rad();
        detector_rot[1] = data.geom.GetPoniRot2_rad();
        double *latt_vec0 = problem.latt_vec0;
        double *latt_vec1 = problem.latt_vec1;
        double *latt_vec2 = problem.latt_vec2;
        double *rot_vec = problem.rot_vec;

        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;
                latt_vec2[1] = 0.0;
                latt_vec2[2] = 0.0;
                break;
            default:
                // Triclinic: initialize a,b,c and α,β,γ from current unit cell
                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
                auto uc = data.latt.GetUnitCell();
                latt_vec2[0] = uc.alpha * PI / 180.0;
                latt_vec2[1] = uc.beta * PI / 180.0;
                latt_vec2[2] = uc.gamma * PI / 180.0;
                break;
        }

        // The spindle. `rocking_spindle` is the fallback for a caller that holds one frame and so
        // passes no `axis` to back-rotate by: the back-rotation is the identity there either way
        // (angle_rad is zero and an AngleAxisRotator of a zero angle-axis ignores the vector, so the
        // block is also held constant), but leaving the {1,0,0} initialiser standing would hand any
        // later reader of this vector the LAB X AXIS in place of the spindle.
        if (const auto spindle = data.axis ? std::optional(data.axis->GetAxis()) : data.rocking_spindle) {
            rot_vec[0] = spindle->x;
            rot_vec[1] = spindle->y;
            rot_vec[2] = spindle->z;
        }

        // The exposure this refinement's spots are spread over, and the spindle they are spread
        // along. Taken from the explicit rocking fields where the caller set them - the per-frame
        // refinement, which does not back-rotate but whose spots still span an exposure - and
        // otherwise from the axis this call does back-rotate by.
        const float rocking_wedge_deg = data.rocking_wedge_deg > 0.0f
                                            ? data.rocking_wedge_deg
                                            : ((data.axis && data.axis->IsScanning())
                                                   ? data.axis->GetWedge_deg() : 0.0f);
        const Coord rocking_spindle = data.rocking_spindle.value_or(
            data.axis ? data.axis->GetAxis() : Coord());
        // Zero everywhere below the trigger, which switches the dead zone off and leaves the gate
        // computing the plain fractional-index miss.
        const double dead_zone_rad = rocking_wedge_deg >= COARSE_SLICING_WEDGE_DEG
                                         ? rocking_wedge_deg * PI / 180.0 * DEAD_ZONE_K
                                         : 0.0;

        const float tolerance_sq = tolerance * tolerance;

        // The same for every spot of every frame, so taken once here rather than per residual.
        const double cos_rot3 = std::cos(data.geom.GetPoniRot3_rad());
        const double sin_rot3 = std::sin(data.geom.GetPoniRot3_rad());

        // Per-image rotation refinement frees only the beam and the orientation and holds the other five
        // blocks constant. Where that is the configuration, the solver uses the reduced residual - the
        // identical fit, with the crystal half worked out once (see XtalResidualBeamOrientation). Any
        // other combination (stills also free the cell, the rotation indexer frees detector angles and
        // spindle) keeps the general form.
        const bool beam_and_orientation_only = data.refine_beam_center
                                               && !data.refine_detector_angles
                                               && !data.refine_rotation_axis
                                               && !data.refine_unit_cell;
        problem.beam_and_orientation_only = beam_and_orientation_only;

        // Sum of w^2 over the spots that entered - the beam prior below is scaled by it so that its
        // strength relative to the data is the same weighted or not. Equals the residual block count
        // when the spots are unweighted.
        double effective_spots = 0.0;

        for (int i = 0; i < spots.size(); i++) {
            if (spots[i].empty())
                continue;

            const std::vector<double> &weight = weights[i]; // empty = unweighted

            double angle_rad = 0.0;
            std::optional<RotMatrix> rot_matr;

            if (data.axis) {
                const float angle_deg = data.axis->GetAngle_deg(i) + data.axis->GetWedge_deg() / 2.0;
                angle_rad = angle_deg * PI / 180.0;
                rot_matr = data.axis->GetTransformationAngle(angle_deg);
            }

            const int frame_index = static_cast<int>(problem.frame_angle_rad.size());
            problem.frame_angle_rad.push_back(angle_rad);

            // Add residuals for each point
            for (size_t j = 0; j < spots[i].size(); j++) {
                const auto &pt = spots[i][j];
                if (!data.index_ice_rings && pt.ice_ring)
                    continue;

                Coord recip = pt.ReciprocalCoord(data.geom);

                if (rot_matr)
                    recip = rot_matr.value() * recip;

                double h_fp = recip * vec0;
                double k_fp = recip * vec1;
                double l_fp = recip * vec2;

                double h = std::round(h_fp);
                double k = std::round(k_fp);
                double l = std::round(l_fp);

                double norm_sq = (h - h_fp) * (h - h_fp) + (k - k_fp) * (k - k_fp) + (l - l_fp) * (l - l_fp);

                // At coarse slicing the spot diffracted somewhere inside the exposure, not at its
                // midpoint, and that unknown angle is a real part of the miss. Charge only the part
                // of it the exposure cannot supply: a rotation delta about the spindle moves the
                // fractional index along u = m x q, so the component of the miss along u is free up
                // to the exposure's rms half-width and only the excess counts. Every other direction
                // is untouched - |q| among them, so every d-spacing is unaffected. Without this the
                // gate is a resolution cut that tightens with the frame width, since the miss grows
                // as a/d.
                if (dead_zone_rad > 0.0) {
                    const Coord u = rocking_spindle % recip;
                    const double u0 = u * vec0, u1 = u * vec1, u2 = u * vec2;
                    const double u_sq = u0 * u0 + u1 * u1 + u2 * u2;
                    if (u_sq > 1e-24) {
                        const double inv_u = 1.0 / std::sqrt(u_sq);
                        const double d_par = ((h - h_fp) * u0 + (k - k_fp) * u1 + (l - l_fp) * u2) * inv_u;
                        const double dead = dead_zone_rad * std::sqrt(u_sq);
                        const double excess = std::max(0.0, std::fabs(d_par) - dead);
                        norm_sq = std::max(0.0, norm_sq - d_par * d_par) + excess * excess;
                    }
                }

                if (norm_sq > tolerance_sq)
                    continue;

                const double weight_sq = weight.empty() ? 1.0 : weight[j] * weight[j];
                effective_spots += weight_sq;

                problem.residuals.emplace_back(pt.x, pt.y,
                                            data.geom.GetWavelength_A(),
                                            data.geom.GetPixelSize_mm(),
                                            cos_rot3, sin_rot3,
                                            angle_rad,
                                            h, k, l,
                                            data.crystal_system,
                                            data.geom.GetOrientation());
                problem.frame.push_back(frame_index);
                // A per-residual weight w enters the squared residual as w^2.
                if (!weight.empty())
                    problem.weight_sq.push_back(weight_sq);
            }
        }

        if (static_cast<int64_t>(problem.residuals.size()) < data.min_spots)
            return false;

        // The gauge direction of a single-axis rotation experiment - parallel to the spindle - written
        // once, for both of the parameter pairs it applies to. The two need it in DIFFERENT frames and
        // that is the whole difficulty:
        //
        //   beam[0]/beam[1] are PIXEL columns and rows. The pixel axes reach the laboratory through
        //   det_matrix = PoniRotMatrix * DetectorOrientation::Matrix(), so on a quarter turn of 1 or 3
        //   the pixel X axis IS the laboratory Y axis. Comparing the goniometer vector's laboratory
        //   components against a beam index is therefore only right when that orientation is the
        //   identity; elsewhere it pins the determined component and frees the gauge one. Project the
        //   spindle onto the pixel axes' own laboratory images instead - exact for any orientation,
        //   any tilt and a spindle at any angle, and equal to picking the dominant component when the
        //   orientation is the identity and the spindle lies along a detector axis.
        //
        //   detector_rot[0]/[1] are rotations about the LABORATORY y and x axes (see PoniRotMatrix),
        //   applied outside that orientation matrix, and they move the direct beam along laboratory x
        //   and y respectively by D/pixel per radian. So the tilt's gauge combination is the spindle's
        //   own laboratory x and y components, with no orientation in it.
        //
        // Same spindle, same physical direction, each in the frame its parameters live in.
        double gauge_beam_x = 0.0, gauge_beam_y = 0.0;
        double gauge_rot_x = 0.0, gauge_rot_y = 0.0;
        if (data.axis) {
            const Coord spindle = data.axis->GetAxis().Normalize();
            const Coord fast = data.geom.GetFastAxis();
            const Coord slow = data.geom.GetSlowAxis();
            const double bx = spindle * fast, by = spindle * slow;
            const double bn = std::hypot(bx, by);
            if (bn > 0.0) {
                gauge_beam_x = bx / bn;
                gauge_beam_y = by / bn;
            }
            const double rn = std::hypot(spindle.x, spindle.y);
            if (rn > 0.0) {
                gauge_rot_x = spindle.x / rn;
                gauge_rot_y = spindle.y / rn;
            }
        }
        // Weight so a gauge prior is a sigma_px-pixel restraint that competes with the positional
        // residuals. k = d|recip|/d(beam_px) ~ pixel/(distance*lambda) [A^-1/px]; scaling by
        // sqrt(#residuals) makes the prior's curvature ~ (1/9) of the well-constrained-data curvature
        // at sigma_px=3, i.e. data wins the perpendicular direction, the prior wins the gauge one.
        // Note what that scaling means: the prior's curvature grows with the number of spots exactly
        // as the data's does, so the split it picks between two aliased parameters is the same however
        // much data the stage has. More frames, a longer sweep or a later stage cannot break it.
        constexpr double sigma_px = 3.0;
        // The tilt's budgets, in those same direct-beam pixels: one for the spindle-parallel
        // combination and one for the perpendicular one. Zero means no restraint at all, so which
        // component is held and which is refined is these two numbers and nothing else.
        //
        // The parallel one is TIGHTER than the beam's on purpose: the data determine the SUM of the
        // two, so with equal budgets the shift splits evenly and half of a beam-centre error still
        // arrives as an angle (measured: the coupling to the starting beam centre falls only from
        // 79% to 41% of one-for-one at equal budgets, and to 8% at this one). The detector tilt is a
        // property of the mounting, re-measured when the detector is calibrated; the beam centre
        // drifts between runs. When both ends of an alias have to be restrained, the tighter
        // restraint belongs on the one that moves less.
        //
        // The perpendicular one is free. That is the arrangement the data support today: it is the
        // component whose conditioning tracks the 2theta the fit reaches, i.e. the one the data speak
        // about, while the parallel one's does not move with 2theta at all.
        constexpr double SIGMA_TILT_PARALLEL_PX = 1.0;
        constexpr double SIGMA_TILT_PERPENDICULAR_PX = 0.0;
        const double gauge_w = data.geom.GetPixelSize_mm() / (distance_mm * data.geom.GetWavelength_A())
                               * std::sqrt(effective_spots) / sigma_px;

        problem.beam_constant = !data.refine_beam_center;
        if (data.refine_beam_center && data.axis) {
            // Gauge handling (single-axis rotation): rotating the whole experiment about the spindle leaves every
            // spot position unchanged, so the beam-centre component PARALLEL to the spindle is a null/gauge-weak
            // direction. Refining it freely lets it wander (~+3 px) and absorb centroid systematics into a wrong
            // beam that the co-refined orientation keeps position-consistent. Rather than freeze it (the beam
            // does drift - it is only LaB6-monitored to ~a few px), RESTRAIN it toward the header with a soft
            // prior: the gauge direction has ~zero data sensitivity so the prior pins it near the header, while a
            // real, well-supported drift can still overcome it.
            problem.priors.push_back({XtalRefinePrior::Block::Beam, gauge_beam_x, gauge_beam_y,
                                      gauge_beam_x * beam[0] + gauge_beam_y * beam[1], gauge_w});
        }

        // Distance, detector angles, rotation axis and cell are parameter blocks only in the general
        // seven-block residual; the reduced one bakes them in, so there is nothing left to configure.
        if (!beam_and_orientation_only) {
            problem.detector_rot_constant = !data.refine_detector_angles;
            if (data.refine_detector_angles) {
                const double rot_range = 3.0 / 180.0 * PI;
                for (int i = 0; i < 2; ++i) {
                    problem.detector_rot_lower[i] = detector_rot[i] - rot_range;
                    problem.detector_rot_upper[i] = detector_rot[i] + rot_range;
                }
                // The same gauge as the beam prior above, described a second time: the tilt moves the
                // direct beam exactly as the beam centre does, at D/pixel px per radian, so leaving
                // its gauge combination free lets a beam-centre error the prior refuses to absorb
                // reappear as an angle - measured at 0.072 deg per pixel of the STARTING beam centre,
                // against a geometric one-for-one of 0.080, while the refined beam never leaves its
                // anchor by more than a quarter of a pixel.
                //
                // Restraining it does not make the tilt a measurement, and nothing here should be read
                // that way. In THIS fit the restrained component carries no information of its own:
                // the crystal orientation is refined alongside it and absorbs the difference, so it
                // ends up as accurate as the file's beam centre and no more. The free component does
                // carry information, and is separately known to sit ~0.06 deg from a powder
                // calibration on one measured detector, which is many times its formal error - so a
                // single crystal's tilt is not a number to feed back into a file. What this buys is
                // that a beam-centre error is no longer laundered into a reported angle.
                //
                // "In this fit" is the load-bearing part: a later stage that FREEZES the orientation
                // has no such compensator, and whether the parallel component is measurable there is a
                // different question with a different answer. This restraint is local to the fit that
                // co-refines the orientation and does not speak for any other.
                if (data.axis) {
                    const double lever = distance_mm / data.geom.GetPixelSize_mm();
                    // Parallel first, then the perpendicular direction (-gy, gx). Both go through the
                    // same restraint, so swapping which one is held is a change to the two budgets.
                    const double dirs[2][2] = {{gauge_rot_x, gauge_rot_y}, {-gauge_rot_y, gauge_rot_x}};
                    const double budget[2] = {SIGMA_TILT_PARALLEL_PX, SIGMA_TILT_PERPENDICULAR_PX};
                    for (int i = 0; i < 2; ++i) {
                        if (budget[i] <= 0.0)
                            continue;
                        problem.priors.push_back({XtalRefinePrior::Block::DetectorRot, dirs[i][0], dirs[i][1],
                                                  dirs[i][0] * detector_rot[0] + dirs[i][1] * detector_rot[1],
                                                  gauge_w * (sigma_px / budget[i]) * lever});
                    }
                }
            }

            problem.rot_vec_constant = !data.refine_rotation_axis;
            if (data.refine_rotation_axis) {
                // Only the DIRECTION of the goniometer axis is a parameter. The residual applies
                // angle_rad * |rot_vec|, so a free three-vector also fits a rotation SCALE - which
                // GoniometerAxis::Axis() then normalises away, leaving the candidate scored by
                // RotationIndexer::accumulate() under a rotation model the fit did not use. Measured
                // over the corpus, that length reached 1.2 % and the fit/score disagreement a whole
                // degree of goniometer angle. It is not a usable measurement either: on synthetic
                // data it recovers 54 % of a known scale error, repeated first passes on one dataset
                // disagree with each other in SIGN, and on the one dataset with a real 1.3 % stage
                // fault it comes out negative. The rotation scale is measured properly, once, with
                // four gates and a jackknife, in PostRefine. Refined on the sphere (see SolveXtalRefine).
            }

            problem.latt_vec1_constant = !data.refine_unit_cell;
            problem.latt_vec2_constant = !data.refine_unit_cell;
            if (data.refine_unit_cell) {
                // Parameter bounds
                // Lengths
                for (int i = 0; i < 3; ++i) {
                    problem.latt_vec1_lower[i] = data.min_length_A;
                    problem.latt_vec1_upper[i] = data.max_length_A;
                }

                if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
                    problem.latt_vec2_constant = false;
                    problem.latt_vec2_lower[0] = std::max(1e-6, PI * (data.min_angle_deg / 180.0));
                    problem.latt_vec2_upper[0] = std::min(PI - 1e-6, PI * (data.max_angle_deg / 180.0));
                } else if (data.crystal_system == gemmi::CrystalSystem::Triclinic) {
                    // α, β, γ bounds (radians)
                    const double alo = PI * (data.min_angle_deg / 180.0);
                    const double ahi = PI * (data.max_angle_deg / 180.0);
                    for (int i = 0; i < 3; ++i) {
                        problem.latt_vec2_lower[i] = alo;
                        problem.latt_vec2_upper[i] = ahi;
                    }
                } else {
                    // Orthorhombic / Tetragonal / Cubic / Hexagonal:
                    // latt_vec2 has no meaning for these systems — always freeze it.
                    problem.latt_vec2_constant = true;
                }
            }
        }

        // Stopping rule: a bound on iterations is reproducible, a bound on wall-clock time is not (see
        // XtalOptimizerData::max_iterations).
        if (data.max_iterations > 0)
            problem.options.max_iterations = data.max_iterations;
        else
            problem.options.max_time_s = data.max_time;
        const LMSummary summary = SolveXtalRefine(problem, num_threads);

        // Only a genuine numerical failure is rejected here: a solve that ran out of iterations or
        // out of time but still descended counts as usable, which is what the real-time caller
        // relies on when it sets max_solver_time. Checked before anything is written back, so a
        // failed refinement leaves data untouched rather than committing half a fit.
        if (!summary.IsSolutionUsable())
            return false;

        if (data.refine_beam_center) {
            data.beam_corr_x = data.geom.GetBeamX_pxl() - beam[0];
            data.beam_corr_y = data.geom.GetBeamY_pxl() - beam[1];
            data.geom.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]);
        }


        if (data.refine_detector_angles)
            data.geom.PoniRot1_rad(detector_rot[0]).PoniRot2_rad(detector_rot[1]);

        if (data.axis && data.refine_rotation_axis)
            data.axis.value().Axis(Coord(rot_vec[0], rot_vec[1], rot_vec[2]));

        if (data.crystal_system == gemmi::CrystalSystem::Orthorhombic)
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
        else if (data.crystal_system == gemmi::CrystalSystem::Tetragonal) {
            latt_vec1[1] = latt_vec1[0];
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
        } else if (data.crystal_system == gemmi::CrystalSystem::Cubic) {
            latt_vec1[1] = latt_vec1[0];
            latt_vec1[2] = latt_vec1[0];
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
        } else if (data.crystal_system == gemmi::CrystalSystem::Hexagonal) {
            latt_vec1[1] = latt_vec1[0];
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1,PI / 2.0, PI / 2.0, 2.0 * PI / 3.0);
        } else if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, latt_vec2[0], PI / 2.0);
        } else {
            // Triclinic via the same generic builder
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, latt_vec2[0], latt_vec2[1], latt_vec2[2]);
        }
        return true;
    } catch (...) {
        // Convergence problems, likely not updated
        return false;
    }
}

bool XtalOptimizer(XtalOptimizerData &data, std::span<const std::vector<SpotToSave>> spots,
                   int num_threads) {
    // A spot's confidence weight is set by its resolution and its intensity, neither of which the solver
    // touches, so the three passes below all get the same weights: take them once.
    std::vector<std::vector<double>> weights(spots.size());
    if (data.weight_spots_by_confidence)
        for (size_t i = 0; i < spots.size(); i++)
            if (!spots[i].empty())
                weights[i] = SpotConfidenceWeights(spots[i]);

    if (!XtalOptimizerInternal(data, spots, weights, XTAL_OPTIMIZER_WIDE_TOLERANCE, num_threads))
        return false;
    XtalOptimizerInternal(data, spots, weights, 0.2, num_threads);
    return XtalOptimizerInternal(data, spots, weights, 0.1, num_threads);
}

bool XtalOptimizer(XtalOptimizerData &data, const std::vector<SpotToSave> &spots, int num_threads) {
    return XtalOptimizer(data, std::span(&spots, 1), num_threads);
}

bool XtalOptimizerRotationOnly(XtalOptimizerData &data,
                               const std::vector<SpotToSave> &spots,
                               const float tolerance) {
    try {
        // Same refusal as XtalOptimizerInternal: the residual here is built from Astar/Bstar/Cstar,
        // which divide by the cell volume, so a coplanar basis makes every one of them infinite.
        if (data.latt.VolumeFraction() < MIN_BASIS_VOLUME_FRACTION)
            return false;

        // Parameter: angle-axis for the extra rotation. Identity == {0,0,0}.
        std::vector<double> rot_aa = {0.0, 0.0, 0.0};

        // Spot selection by current indexing (same approach as XtalOptimizerInternal)
        const Coord a0 = data.latt.Vec0();
        const Coord b0 = data.latt.Vec1();
        const Coord c0 = data.latt.Vec2();

        const float tol_sq = tolerance * tolerance;

        // Each selected spot: its observed reciprocal vector and the indices it is fitted to.
        struct Observation {
            Coord s_obs;
            double h, k, l;
        };
        std::vector<Observation> observations;

        for (const auto &pt : spots) {
            if (!data.index_ice_rings && pt.ice_ring)
                continue;

            // Compute fractional HKL using the CURRENT lattice
            Coord recip_index = pt.ReciprocalCoord(data.geom);
            if (data.axis.has_value())
                recip_index = data.axis->GetTransformationAngle(pt.phi) * recip_index;

            const double h_fp = static_cast<double>(recip_index * a0);
            const double k_fp = static_cast<double>(recip_index * b0);
            const double l_fp = static_cast<double>(recip_index * c0);

            const double h = std::round(h_fp);
            const double k = std::round(k_fp);
            const double l = std::round(l_fp);

            const double norm_sq =
                (h - h_fp) * (h - h_fp) +
                (k - k_fp) * (k - k_fp) +
                (l - l_fp) * (l - l_fp);

            if (norm_sq > static_cast<double>(tol_sq))
                continue;

            // s_obs must be in the same reference frame as the
            // predicted reciprocal vector (h·a* + k·b* + l·c*), which is the
            // phi=0 crystal frame.  Apply the same goniometer back-rotation
            // that was used above for the HKL assignment.
            Coord s_obs = data.geom.DetectorToRecip(pt.x, pt.y);
            if (data.axis.has_value())
                s_obs = data.axis->GetTransformationAngle(pt.phi) * s_obs;

            observations.push_back({s_obs, h, k, l});
        }

        if (static_cast<int64_t>(observations.size()) < data.min_spots)
            return false;

        // Residual: s_obs - R(rot_aa) (h a* + k b* + l c*), the reciprocal basis rotated once per
        // evaluation and shared by every spot.
        //
        // Regularization: prefer the smallest rotation correction that fits the data, w * rot_aa. This is
        // essential when spots are nearly coplanar in reciprocal space (e.g. still images), where the
        // rotation component perpendicular to the scattering plane is otherwise underdetermined. The
        // weight is in A^-1 rad^-1, relative to the typical residual.
        const double reg_weight = 0.05;
        const Coord astar = data.latt.Astar(), bstar = data.latt.Bstar(), cstar = data.latt.Cstar();
        const auto evaluate = [&](const double *aa, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) {
            using D = Dual<3>;
            const D aa_d[3] = {D::Variable(aa[0], 0), D::Variable(aa[1], 1), D::Variable(aa[2], 2)};
            const AngleAxisRotator<D> rot(aa_d);
            const double astar_unrot[3] = {astar.x, astar.y, astar.z};
            const double bstar_unrot[3] = {bstar.x, bstar.y, bstar.z};
            const double cstar_unrot[3] = {cstar.x, cstar.y, cstar.z};
            D astar_rot[3], bstar_rot[3], cstar_rot[3];
            rot.Rotate(astar_unrot, astar_rot);
            rot.Rotate(bstar_unrot, bstar_rot);
            rot.Rotate(cstar_unrot, cstar_rot);

            cost = 0.0;
            const auto add = [&](double r, const double *J) {
                cost += 0.5 * r * r;
                if (!g)
                    return;
                for (int i = 0; i < 3; i++) {
                    (*g)[i] += J[i] * r;
                    for (int j = 0; j < 3; j++)
                        (*H)(i, j) += J[i] * J[j];
                }
            };
            for (const auto &o: observations) {
                const double s_obs[3] = {o.s_obs.x, o.s_obs.y, o.s_obs.z};
                for (int c = 0; c < 3; c++) {
                    const D pred = o.h * astar_rot[c] + o.k * bstar_rot[c] + o.l * cstar_rot[c];
                    const double J[3] = {-pred.v[0], -pred.v[1], -pred.v[2]};
                    add(s_obs[c] - pred.a, J);
                }
            }
            for (int c = 0; c < 3; c++) {
                double J[3] = {0.0, 0.0, 0.0};
                J[c] = reg_weight;
                add(reg_weight * aa[c], J);
            }
            return std::isfinite(cost) && (!g || (g->allFinite() && H->allFinite()));
        };

        std::vector<LMBlock> blocks(1);
        blocks[0].size = 3;
        LMOptions options;
        if (data.max_iterations > 0)
            options.max_iterations = data.max_iterations;
        else
            options.max_time_s = data.max_time;
        const LMSummary summary = SolveLM(rot_aa, blocks, options, evaluate);

        if (!summary.IsSolutionUsable())
            return false;

        // Apply rotation to direct-lattice vectors.
        // ceres::AngleAxisToRotationMatrix writes a **row-major** 3×3 matrix,
        // and Eigen's << operator also fills row-by-row, so the assignment
        // below is correct without any transposing.
        //
        // Note: for a pure orthogonal rotation R, R⁻ᵀ = R, so rotating the
        // direct-lattice vectors (A, B, C) by R is exactly equivalent to
        // rotating the reciprocal vectors (a*, b*, c*) by the same R.  No
        // transpose or inversion of R is needed here.
        double R_raw[9];
        ceres::AngleAxisToRotationMatrix(rot_aa.data(), R_raw); // row-major 3x3

        Eigen::Matrix3d R;
        R << R_raw[0], R_raw[3], R_raw[6],
             R_raw[1], R_raw[4], R_raw[7],
             R_raw[2], R_raw[5], R_raw[8];

        const Eigen::Vector3d A(a0.x, a0.y, a0.z);
        const Eigen::Vector3d B(b0.x, b0.y, b0.z);
        const Eigen::Vector3d C(c0.x, c0.y, c0.z);

        const Eigen::Vector3d A2 = R * A;
        const Eigen::Vector3d B2 = R * B;
        const Eigen::Vector3d C2 = R * C;

        data.latt = CrystalLattice(
            Coord(static_cast<float>(A2.x()), static_cast<float>(A2.y()), static_cast<float>(A2.z())),
            Coord(static_cast<float>(B2.x()), static_cast<float>(B2.y()), static_cast<float>(B2.z())),
            Coord(static_cast<float>(C2.x()), static_cast<float>(C2.y()), static_cast<float>(C2.z()))
        );

        double theta = std::sqrt(rot_aa[0] * rot_aa[0] + rot_aa[1] * rot_aa[1] + rot_aa[2] * rot_aa[2]);
        data.angle_corr = theta;
        if (theta > 1e-6) {
            Coord rot;
            rot.x = rot_aa[0] / theta;
            rot.y = rot_aa[1] / theta;
            rot.z = rot_aa[2] / theta;
            data.angle_axis = rot;
        } else
            data.angle_axis.reset();

        return true;
    } catch (...) {
        return false;
    }
}