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

#include "StillsPartialityRefine.h"

#include <algorithm>
#include <array>
#include <atomic>
#include <cmath>
#include <future>
#include <thread>
#include <vector>

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

#include "Merge.h"
#include "../../common/Logger.h"
#include "../../common/ThreadAffinity.h"

namespace {
    constexpr size_t MIN_FIT_REFLECTIONS = 20;
    constexpr double kRadToDeg = 180.0 / 3.14159265358979323846;

    // Model error of the per-crystal scale fit, as a fraction of the predicted intensity - ScaleOnTheFly's.
    constexpr double SCALE_MODEL_ERROR = 0.3;

    double SafeInv(double x, double fallback) {
        if (!std::isfinite(x) || x == 0.0)
            return fallback;
        return 1.0 / x;
    }

    // One accepted reflection reduced to the physical partiality fit: the base reciprocal vector q (crystal
    // frame at the stored per-image orientation), the reference full intensity, the measured intensity, the
    // reciprocal of the prescaling correction (the factor that turns a true intensity back into counts),
    // the measured variance and its background part, and the weight the tilt is refined with.
    // dist_ewald / partiality are recomputed from q under the refined tilt.
    struct FitObs {
        double qx, qy, qz;
        double Iref;
        double Iobs;
        double inv_corr;     // 1 / (prescaling_corr * qe_corr * flight_corr)
        double sigma2;
        double var_bkg;
        double weight;       // 1 / sqrt(variance at the current model), set with each G
    };

    // The variance of an observation at the model m: background, the counting noise of m and the model error.
    double ModelVariance(const FitObs &o, double m) {
        return o.var_bkg + std::max(0.0, m) + SCALE_MODEL_ERROR * SCALE_MODEL_ERROR * m * m;
    }

    double VecLen(double x, double y, double z) { return std::sqrt(x * x + y * y + z * z); }

    // The rocking width sigma^2 at a reflection of resolution d* (and |q_z|, for the bandwidth term): the
    // run's table where there is one - linear in d* between its shells, held beyond the outermost - else
    // the three-term formula.
    struct Width {
        double gamma0, gamma_e, bw;
        const std::vector<double> *table_dstar = nullptr, *table_s = nullptr;

        double Sigma2(double dstar, double qz) const {
            if (table_dstar && !table_dstar->empty()) {
                const auto &x = *table_dstar;
                const auto &y = *table_s;
                double s = y.front();
                if (dstar >= x.back())
                    s = y.back();
                else if (dstar > x.front()) {
                    const size_t i = std::upper_bound(x.begin(), x.end(), dstar) - x.begin();
                    s = y[i - 1] + (dstar - x[i - 1]) / (x[i] - x[i - 1]) * (y[i] - y[i - 1]);
                }
                return s * s;
            }
            const double sig_ang = gamma_e * dstar;
            const double sbw = bw * std::fabs(qz);
            return gamma0 * gamma0 + sig_ang * sig_ang + sbw * sbw;
        }
    };

    // Analytic partiality for a reflection whose base reciprocal vector is q, tilted by (psi_x, psi_y).
    // Mirrors BraggPrediction: dist_ewald = |S| - 1/lambda with S = q_rot + S0, and
    // p = exp(-dist_ewald^2 / 2 sigma^2), sigma^2 = gamma0^2 + (gamma_e*d*)^2 + (bw*|q_z|)^2 - three
    // independent broadenings added in quadrature: the reciprocal-lattice point's own radius (gamma0,
    // ~1/domain size, resolution-INdependent), the mosaic/divergence spread (gamma_e*d*, proportional to
    // d*) and the bandwidth smear along the beam. In practice the fit below returns gamma_e ~ 0 and the
    // width is essentially gamma0 - see there.
    double ComputeP(double qx, double qy, double qz,
                    double psi_x, double psi_y,
                    double s0x, double s0y, double s0z, double inv_lambda,
                    const Width &width) {
        double aa[3] = {psi_x, psi_y, 0.0};
        double q[3] = {qx, qy, qz};
        double qr[3];
        ceres::AngleAxisRotatePoint(aa, q, qr);
        const double Sx = qr[0] + s0x, Sy = qr[1] + s0y, Sz = qr[2] + s0z;
        const double de = std::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - inv_lambda;
        const double dstar = std::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]);
        const double sig2 = width.Sigma2(dstar, qr[2]);
        if (!(sig2 > 0.0))
            return 1.0;
        return std::exp(-0.5 * de * de / sig2);
    }

    // The prediction window cuts the rocking curve. Reflections fill a window of half-width W evenly, so
    // their plain moment of de is u = W^2 / 3, and the intensity-weighted moment m of a Gaussian of width
    // s cut at +-W is u * g(a), a = W / s, with g(a) = 3/a^2 (1 - 2 a phi(a) / erf(a / sqrt 2)): 1 for a
    // curve much wider than the window, s^2 / u for a much narrower one. Solves g(a) = m / u for s^2 by
    // bisection (g falls with a). A curve wider than the window cannot be measured through it: empty.
    std::optional<double> UntruncatedVariance(double m, double u) {
        auto g = [](double a) {
            const double phi = std::exp(-0.5 * a * a) / std::sqrt(2.0 * 3.14159265358979323846);
            return 3.0 / (a * a) * (1.0 - 2.0 * a * phi / std::erf(a / std::sqrt(2.0)));
        };
        const double t = m / u;
        double lo = 1.0, hi = 50.0;
        if (!(t < g(lo)))
            return std::nullopt;
        for (int i = 0; i < 60; ++i) {
            const double mid = 0.5 * (lo + hi);
            (g(mid) > t ? lo : hi) = mid;
        }
        const double a = 0.5 * (lo + hi);
        return 3.0 * u / (a * a);
    }

    // Robust per-crystal scale G (linear in G given the model coefficients), the objective of
    // ScaleOnTheFly::SolveScaleIRLS: minimise sum Cauchy_k( (G*coeff - Iobs) / sqrt(var(G*coeff)) ) over G >= 0,
    // the variances re-evaluated at each step's G.
    double SolveScaleIRLS(const std::vector<FitObs> &obs, const std::vector<double> &coeff, double robust_k) {
        double num = 0.0, den = 0.0;
        for (size_t i = 0; i < obs.size(); ++i) {
            num += coeff[i] * obs[i].Iobs / obs[i].sigma2;
            den += coeff[i] * coeff[i] / obs[i].sigma2;
        }
        if (!(den > 0.0) || !std::isfinite(num / den))
            return 1.0;
        double G = std::max(0.0, num / den);

        const double k2 = robust_k * robust_k;
        for (int iter = 0; iter < 30; ++iter) {
            num = den = 0.0;
            for (size_t i = 0; i < obs.size(); ++i) {
                const double res = G * coeff[i] - obs[i].Iobs;
                const double w = 1.0 / (ModelVariance(obs[i], G * coeff[i]) + res * res / k2);
                num += w * coeff[i] * obs[i].Iobs;
                den += w * coeff[i] * coeff[i];
            }
            if (!(den > 0.0) || !std::isfinite(num / den))
                break;
            const double G_prev = G;
            G = std::max(0.0, num / den);
            if (std::abs(G - G_prev) <= 1e-7 * std::max(G, 1.0))
                break;
        }
        return G;
    }

    // The rotation by the angle-axis vector (psi_x, psi_y, 0) as a matrix (Rodrigues), so that one tilt
    // is applied to all the reflections of a crystal by a matrix product each.
    std::array<double, 9> TiltMatrix(double psi_x, double psi_y) {
        const double theta = std::sqrt(psi_x * psi_x + psi_y * psi_y);
        if (theta < 1e-15)
            return {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0};
        const double kx = psi_x / theta, ky = psi_y / theta;
        const double c = std::cos(theta), s = std::sin(theta), t = 1.0 - c;
        return {c + t * kx * kx, t * kx * ky, s * ky,
                t * kx * ky, c + t * ky * ky, -s * kx,
                -s * ky, s * kx, c};
    }

    // ComputeP with the tilt given as its rotation matrix.
    struct TiltModel {
        double s0x, s0y, s0z, inv_lambda;
        Width width;

        double P(const std::array<double, 9> &R, const FitObs &o) const {
            const double qx = R[0] * o.qx + R[1] * o.qy + R[2] * o.qz;
            const double qy = R[3] * o.qx + R[4] * o.qy + R[5] * o.qz;
            const double qz = R[6] * o.qx + R[7] * o.qy + R[8] * o.qz;
            const double de = VecLen(qx + s0x, qy + s0y, qz + s0z) - inv_lambda;
            const double sig2 = width.Sigma2(VecLen(qx, qy, qz), qz);
            if (!(sig2 > 0.0))
                return 1.0;
            return std::exp(-0.5 * de * de / sig2);
        }
    };

    // Refine the orientation tilt (psi_x, psi_y) holding the scale G fixed: minimise
    //   sum_i rho(r_i^2) / 2 + |psi / sigma_prior|^2 / 2,   r_i = weight_i (G p_i(psi) Iref_i / corr_i - Iobs_i),
    // with the Cauchy rho(s) = k^2 log(1 + s / k^2) - the intensity-space objective of ScaleOnTheFly, and
    // the Gaussian prior a proper MAP term since r_i is already in sigma units. Levenberg-Marquardt on the
    // normal equations reweighted by rho', the Jacobian by central differences, every trial tilt held inside
    // the bound. Returns false when the cost cannot be evaluated, leaving psi as it came in.
    bool RefineTilt(const std::vector<FitObs> &obs, const TiltModel &model, double G, double robust_k,
                    double inv_sigma_prior, double max_tilt, double psi[2]) {
        constexpr int MAX_ITERATIONS = 25;
        constexpr double STEP = 1e-7;   // rad, for the central differences
        const double k2 = robust_k * robust_k;
        const double prior2 = inv_sigma_prior * inv_sigma_prior;

        auto cost = [&](const double p[2]) {
            const auto R = TiltMatrix(p[0], p[1]);
            double c = 0.5 * prior2 * (p[0] * p[0] + p[1] * p[1]);
            for (const auto &o: obs) {
                const double r = o.weight * (G * model.P(R, o) * o.inv_corr * o.Iref - o.Iobs);
                c += 0.5 * k2 * std::log1p(r * r / k2);
            }
            return c;
        };

        double current = cost(psi);
        if (!std::isfinite(current))
            return false;
        double lambda = 1e-4;
        for (int iter = 0; iter < MAX_ITERATIONS; ++iter) {
            const auto R = TiltMatrix(psi[0], psi[1]);
            const auto Rxp = TiltMatrix(psi[0] + STEP, psi[1]), Rxm = TiltMatrix(psi[0] - STEP, psi[1]);
            const auto Ryp = TiltMatrix(psi[0], psi[1] + STEP), Rym = TiltMatrix(psi[0], psi[1] - STEP);
            double A00 = prior2, A01 = 0.0, A11 = prior2;
            double b0 = prior2 * psi[0], b1 = prior2 * psi[1];
            for (const auto &o: obs) {
                const double a = o.weight * G * o.inv_corr * o.Iref;
                const double r = a * model.P(R, o) - o.weight * o.Iobs;
                const double jx = a * (model.P(Rxp, o) - model.P(Rxm, o)) / (2.0 * STEP);
                const double jy = a * (model.P(Ryp, o) - model.P(Rym, o)) / (2.0 * STEP);
                const double w = 1.0 / (1.0 + r * r / k2);   // rho'
                A00 += w * jx * jx;
                A01 += w * jx * jy;
                A11 += w * jy * jy;
                b0 += w * jx * r;
                b1 += w * jy * r;
            }

            bool accepted = false;
            double trial[2] = {psi[0], psi[1]};
            double trial_cost = current;
            while (lambda < 1e10) {
                const double M00 = A00 * (1.0 + lambda), M11 = A11 * (1.0 + lambda);
                const double det = M00 * M11 - A01 * A01;
                if (!(det > 0.0))
                    break;
                trial[0] = std::clamp(psi[0] - (M11 * b0 - A01 * b1) / det, -max_tilt, max_tilt);
                trial[1] = std::clamp(psi[1] - (M00 * b1 - A01 * b0) / det, -max_tilt, max_tilt);
                trial_cost = cost(trial);
                if (std::isfinite(trial_cost) && trial_cost < current) {
                    accepted = true;
                    lambda = std::max(lambda / 10.0, 1e-12);
                    break;
                }
                lambda *= 10.0;
            }
            if (!accepted)
                break;
            const double step = std::hypot(trial[0] - psi[0], trial[1] - psi[1]);
            const double decrease = current - trial_cost;
            psi[0] = trial[0];
            psi[1] = trial[1];
            current = trial_cost;
            if (decrease <= 1e-6 * current || step <= 1e-10)
                break;
        }
        return true;
    }
}

StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x)
    : experiment_(x),
      hkl_key_generator_(x.GetScalingSettings().GetMergeFriedel(), x.GetSpaceGroupOrP1()),
      d_min_limit_(x.GetScalingSettings().GetHighResolutionLimit_A()),
      d_max_limit_(x.GetScalingSettings().GetLowResolutionLimit_A()),
      min_partiality_(x.GetScalingSettings().GetMinPartiality()),
      bandwidth_sigma_(x.GetBandwidthFWHM().value_or(0.0f) / 2.3548f) {}

double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome,
                                         const std::map<HKLKey, double> &reference,
                                         const std::optional<RunWidth> &run_width,
                                         double radial_scale) const {
    if (outcome.reflections.empty())
        return 0.0;

    const Coord Astar = outcome.latt.Astar();
    const Coord Bstar = outcome.latt.Bstar();
    const Coord Cstar = outcome.latt.Cstar();
    const Coord S0 = outcome.geom.GetScatteringVector();
    const double inv_lambda = 1.0 / outcome.geom.GetWavelength_A();

    auto base_q = [&](const Reflection &r) {
        return (Astar * static_cast<float>(r.h) + Bstar * static_cast<float>(r.k)
                + Cstar * static_cast<float>(r.l)) * static_cast<float>(radial_scale);
    };

    // Collect the reflections that constrain the fit (accepted, non-ice, finite, present in the reference).
    std::vector<FitObs> obs;
    obs.reserve(outcome.reflections.size());
    // Moments of the excitation error against resolution: de^2 ~ gamma0^2 + gamma_e^2 * d*^2, fitted per
    // crystal by ordinary least squares on (d*^2, de^2). Both components come out of the data.
    double m_n = 0.0, m_x = 0.0, m_xx = 0.0, m_y = 0.0, m_xy = 0.0;
    size_t n_de = 0;
    for (const Reflection &r: outcome.reflections) {
        if (r.on_ice_ring || !AcceptReflection(r, d_min_limit_, d_max_limit_))
            continue;
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f)
            continue;
        const auto it = reference.find(hkl_key_generator_(r));
        if (it == reference.end() || !std::isfinite(it->second))
            continue;

        const Coord q = base_q(r);
        obs.push_back(FitObs{
            .qx = q.x, .qy = q.y, .qz = q.z,
            .Iref = it->second,
            .Iobs = static_cast<double>(r.I),
            .inv_corr = SafeInv(r.prescaling_corr * r.qe_corr * r.flight_corr, 1.0),
            .sigma2 = static_cast<double>(r.sigma) * r.sigma,
            .var_bkg = std::isfinite(r.var_bkg) && r.var_bkg > 0.0f ? r.var_bkg : static_cast<double>(r.sigma) * r.sigma,
            .weight = SafeInv(r.sigma, 1.0),
        });

        // Excitation error at the stored orientation (psi = 0), collected as the moments of de^2 against
        // d*^2, so BOTH width components are fitted rather than one being forced to zero. Forcing the
        // width to be purely angular (gamma0 = 0) pins it to the high-resolution edge - it is fitted over
        // a d*^2-dense population - and it then collapses at low d*, giving p ~ 0 for reflections that
        // were plainly recorded, which inflated the merged low-resolution intensity scale ~3.6x.
        const double dstar = VecLen(q.x, q.y, q.z);
        const double de0 = VecLen(q.x + S0.x, q.y + S0.y, q.z + S0.z) - inv_lambda;
        if (dstar > 1e-9) {
            const double x = dstar * dstar, y = de0 * de0;
            m_n += 1.0;
            m_x += x;
            m_xx += x * x;
            m_y += y;
            m_xy += x * y;
            ++n_de;
        }
    }

    if (obs.size() < MIN_FIT_REFLECTIONS || n_de == 0)
        return 0.0;

    // Without the run's measured width (too few reflections in the run to measure it), the width comes
    // from this crystal's own reflections. Solve the 2x2 normal equations for de^2 = A + B d*^2. A degenerate spread in d* (all reflections in
    // one shell) leaves B undetermined, so fall back to the pure angular width there; a negative fitted
    // component is unphysical and is clamped to zero, which reduces to the previous model.
    //
    // Measured outcome, worth knowing before touching this: the fit does NOT split the width between the
    // two terms - it returns gamma0 ~ 4e-4 1/A and gamma_e ~ 0 (their cross-over sits at d = 0.3 A, far
    // outside any measured range), i.e. a width constant in the LINEAR Ewald distance. That is
    // structural, not a fluke: prediction accepts reflections on a fixed linear |dist_ewald| cutoff, so
    // the accepted population's de^2 is flat in d*^2 by construction and the slope is genuinely ~0. The
    // truncated population cannot constrain an angular term, and its width is the window's, not the
    // crystal's - which is why the run's measured width is used wherever there is one.
    const double det = m_n * m_xx - m_x * m_x;
    double A = 0.0, B = 0.0;
    if (std::fabs(det) > 1e-30) {
        A = (m_xx * m_y - m_x * m_xy) / det;
        B = (m_n * m_xy - m_x * m_y) / det;
    } else {
        B = m_x > 0.0 ? m_y / m_x : 0.0;
    }
    double gamma0 = std::sqrt(std::max(0.0, A));
    const double gamma_e_fit = std::max(std::sqrt(std::max(0.0, B)), 1e-9);
    double gamma_e = settings_.gamma_e > 0.0 ? settings_.gamma_e : gamma_e_fit;
    double bw = bandwidth_sigma_;
    if (run_width) {
        gamma0 = run_width->gamma0;
        gamma_e = std::max(run_width->gamma_e, 1e-9);
        bw = run_width->bandwidth_sigma;
    }
    const Width width{.gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw,
                      .table_dstar = run_width ? &run_width->table_dstar : nullptr,
                      .table_s = run_width ? &run_width->table_s : nullptr};

    double psi[2] = {0.0, 0.0};
    double G = 1.0;

    const TiltModel model{.s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda,
                          .width = width};
    const bool refine_tilt = obs.size() >= settings_.min_reflections;
    const int inner = refine_tilt ? settings_.inner_iterations : 1;

    for (int it = 0; it < inner; ++it) {
        // (1) Solve G given the current partialities, and weight the tilt fit by the variances at that model.
        std::vector<double> coeff(obs.size());
        for (size_t j = 0; j < obs.size(); ++j) {
            const double p = ComputeP(obs[j].qx, obs[j].qy, obs[j].qz, psi[0], psi[1],
                                      S0.x, S0.y, S0.z, inv_lambda, width);
            coeff[j] = p * obs[j].inv_corr * obs[j].Iref;
        }
        G = SolveScaleIRLS(obs, coeff, settings_.robust_k);
        if (!(G > 0.0) || !std::isfinite(G))
            return 0.0;
        for (size_t j = 0; j < obs.size(); ++j)
            obs[j].weight = 1.0 / std::sqrt(ModelVariance(obs[j], G * coeff[j]));

        if (!refine_tilt)
            break;

        // (2) Refine the tilt holding G fixed. A cost that cannot be evaluated keeps the tilt this crystal
        // came in with and stops refining it - G alone is still a usable model.
        const double inv_sigma_prior = settings_.prior_sigma_deg > 0.0 ? kRadToDeg / settings_.prior_sigma_deg : 0.0;
        if (!RefineTilt(obs, model, G, settings_.robust_k, inv_sigma_prior, settings_.max_tilt_rad, psi))
            break;
    }

    // Write the refined partiality + scale correction onto every reflection of the crystal (not only the
    // fit subset), so the merge sees a consistent model. image_scale_corr = prescaling_corr * qe_corr * flight_corr /
    // (partiality * G), the same composition ScaleOnTheFly writes.
    for (auto &r: outcome.reflections) {
        const Coord q = base_q(r);
        const double p = ComputeP(q.x, q.y, q.z, psi[0], psi[1], S0.x, S0.y, S0.z, inv_lambda, width);
        r.partiality = static_cast<float>(p);
        const float corr = r.prescaling_corr * r.qe_corr * r.flight_corr;
        const double denom = p * G;
        r.image_scale_corr = (std::isfinite(corr) && std::isfinite(denom) && denom > 0.0)
                                 ? static_cast<float>(corr / denom)
                                 : NAN;
    }
    outcome.image_scale_g = static_cast<float>(G);

    // The corrections just changed, so the CC that ScaleOnTheFly measured no longer describes them.
    // Refresh it here: it is reported per image and --min-image-cc drops images by it, so it has to be
    // the CC of the data that is actually merged.
    const auto [cc, cc_n] = ImageReferenceCC(outcome.reflections, reference, hkl_key_generator_,
                                             d_min_limit_, d_max_limit_, min_partiality_);

    outcome.image_scale_cc = cc;
    outcome.image_scale_cc_n = cc_n;

    const double tilt_deg = std::sqrt(psi[0] * psi[0] + psi[1] * psi[1]) * kRadToDeg;
    return tilt_deg;
}

double StillsPartialityRefine::MeasureRadialScale(const std::vector<IntegrationOutcome> &outcomes,
                                                  const std::optional<UnitCell> &reference_cell,
                                                  double radial_scale) const {
    // A scale of the reciprocal lattice - a wavelength, or a cell, a little off - leaves every spot where
    // it was (positions fix only lambda d*), but moves each reflection's excitation error by
    // eps lambda d*^2 / 2 for q -> (1 + eps) q, so the run's rocking curve sits off zero by an amount
    // growing as d*^2. The prediction window is symmetric in de, so the intensity-weighted mean de of a
    // shell is that offset, and a line through the origin in lambda d*^2 / 2 gives eps. A window that cuts
    // an off-centre curve pulls the mean toward zero, so the caller iterates from the scale found.
    const auto &indexing = experiment_.GetIndexingSettings();
    struct Point { double x, de, w; };
    std::vector<Point> pts;
    double max_x = 0.0;
    for (const auto &o: outcomes) {
        if (!o.image_scale_g || !(*o.image_scale_g > 0.0))
            continue;
        if (reference_cell && !o.latt.GetUnitCell().is_close(*reference_cell, indexing.GetUnitCellDistTolerance(),
                                                               indexing.GetUnitCellAngleTolerance_deg()))
            continue;
        const Coord S0 = o.geom.GetScatteringVector();
        const double inv_lambda = 1.0 / o.geom.GetWavelength_A();
        for (const auto &r: o.reflections) {
            if (r.on_ice_ring || !AcceptReflection(r, d_min_limit_, d_max_limit_) || !(r.d > 0.0f))
                continue;
            const double w = r.I * r.prescaling_corr * r.qe_corr * r.flight_corr / *o.image_scale_g;
            if (!std::isfinite(w))
                continue;
            const Coord q = (o.latt.Astar() * static_cast<float>(r.h) + o.latt.Bstar() * static_cast<float>(r.k)
                             + o.latt.Cstar() * static_cast<float>(r.l)) * static_cast<float>(radial_scale);
            const double x = 0.5 / (inv_lambda * r.d * r.d);
            pts.push_back({x, VecLen(q.x + S0.x, q.y + S0.y, q.z + S0.z) - inv_lambda, w});
            max_x = std::max(max_x, x);
        }
    }
    constexpr int N_SHELLS = 8;
    if (pts.size() < 100 * N_SHELLS)
        return radial_scale;
    std::array<double, N_SHELLS> sw{}, swde{}, swx{};
    for (const auto &p: pts) {
        const int b = std::min(N_SHELLS - 1, static_cast<int>(p.x / max_x * N_SHELLS));
        sw[b] += p.w;
        swde[b] += p.w * p.de;
        swx[b] += p.w * p.x;
    }
    double sxy = 0.0, sxx = 0.0;
    for (int b = 0; b < N_SHELLS; ++b) {
        if (!(sw[b] > 0.0))
            continue;
        const double x = swx[b] / sw[b], y = swde[b] / sw[b];
        sxy += sw[b] * x * y;
        sxx += sw[b] * x * x;
    }
    if (!(sxx > 0.0) || !std::isfinite(sxy))
        return radial_scale;
    // the slope is eps, removed by scaling q by (1 - eps)
    return radial_scale * (1.0 - sxy / sxx);
}

std::optional<StillsPartialityRefine::RunWidth> StillsPartialityRefine::MeasureRunWidth(
        const std::vector<IntegrationOutcome> &outcomes, const std::optional<UnitCell> &reference_cell,
        double radial_scale, const RunWidth *window_width) const {
    // The rocking width measured on the crystals rather than on the prediction window. A lattice point
    // sits anywhere across the window with equal likelihood, while what it records is I_full * p(de), so
    // the intensity-weighted second moment of the excitation error, sum I de^2 / sum I, is the width s^2
    // of a Gaussian p - corrected for the window cutting the curve's tails (UntruncatedVariance). The
    // plain moment of de over the predicted reflections, which the per-crystal fit in RefineOne uses, is
    // instead the window's own (W^2 / 3). Measured in shells of d*^2 over the whole run, where the
    // intensities of thousands of crystals average out, and fitted as s^2 = gamma0^2 + gamma_e^2 d*^2 +
    // (bandwidth_sigma lambda d*^2 / 2)^2.
    //
    // Each frame predicts through a window of its own (twice its profile radius), and these differ by an
    // order of magnitude within a run. Pooled, the frames with the widest windows decide both moments, and
    // the correction then reads their window: on a run whose windows are mostly narrower than the curve at
    // high resolution, the pooled measurement came out as the window at every resolution, a width constant
    // in d* that is several times too wide at low resolution. So within a shell the frames are taken in
    // WINDOW_GROUPS groups of similar window, each measured against its own, and the shell's width is the
    // median of the groups whose curve is narrower than their window; a shell where no group can see the
    // curve is left out.
    //
    // A frame's window is the largest excitation error it predicted where it predicted through one width.
    // Predicted through the run's width (window_width), the window is its formula at the reflection's own
    // resolution: the frame's largest excitation error is then the high-resolution window, and taken as the
    // window of its low-resolution reflections it put them in the wrong group.
    const auto &indexing = experiment_.GetIndexingSettings();
    struct Point { double x, de2, w, var_w, window; };
    std::vector<Point> pts;
    for (const auto &o: outcomes) {
        if (!o.image_scale_g || !(*o.image_scale_g > 0.0))
            continue;
        if (reference_cell && !o.latt.GetUnitCell().is_close(*reference_cell, indexing.GetUnitCellDistTolerance(),
                                                               indexing.GetUnitCellAngleTolerance_deg()))
            continue;
        float frame_window = 0.0f;   // the largest excitation error the frame predicted
        for (const auto &r: o.reflections)
            if (std::isfinite(r.dist_ewald))
                frame_window = std::max(frame_window, std::fabs(r.dist_ewald));
        for (const auto &r: o.reflections) {
            if (r.on_ice_ring || !AcceptReflection(r, d_min_limit_, d_max_limit_))
                continue;
            const double corr = r.prescaling_corr * r.qe_corr * r.flight_corr / *o.image_scale_g;
            const double w = r.I * corr, sigma_w = r.sigma * corr;
            if (!std::isfinite(w) || !std::isfinite(sigma_w) || !std::isfinite(r.dist_ewald) || !(r.d > 0.0f))
                continue;
            // The excitation error the window was cut in is the prediction's; the curve's is taken at the
            // run's radial scale (MeasureRadialScale), recomputed from the lattice when that is not one.
            double de = r.dist_ewald;
            if (radial_scale != 1.0) {
                const Coord q = (o.latt.Astar() * static_cast<float>(r.h) + o.latt.Bstar() * static_cast<float>(r.k)
                                 + o.latt.Cstar() * static_cast<float>(r.l)) * static_cast<float>(radial_scale);
                const Coord S0 = o.geom.GetScatteringVector();
                de = VecLen(q.x + S0.x, q.y + S0.y, q.z + S0.z) - 1.0 / o.geom.GetWavelength_A();
            }
            double window = frame_window;
            if (window_width) {
                // the formula at the reflection's d*, its bandwidth term at |q_z| = lambda d*^2 / 2
                const double dstar = 1.0 / r.d;
                const double a = window_width->gamma_e * dstar;
                const double b = window_width->bandwidth_sigma * o.geom.GetWavelength_A() * dstar * dstar / 2.0;
                window = std::sqrt(window_width->gamma0 * window_width->gamma0 + a * a + b * b);
            }
            pts.push_back({1.0 / (static_cast<double>(r.d) * r.d), de * de, w, sigma_w * sigma_w, window});
        }
    }
    constexpr size_t N_SHELLS = 8;
    constexpr size_t WINDOW_GROUPS = 4;
    if (pts.size() < 100 * N_SHELLS)
        return std::nullopt;
    std::sort(pts.begin(), pts.end(), [](const Point &a, const Point &b) { return a.x < b.x; });

    Logger logger("StillsPartialityRefine");
    // A group's moment is no better than the summed intensity it is weighted by, so a group is used only
    // where that sum is measured to 10 % (sum I / sigma(sum I) >= 10): beyond the diffraction limit the
    // sum is noise and can even be negative, and says nothing about the width.
    constexpr double MIN_SHELL_SIGNAL = 10.0;
    // The width of the points [lo, hi) of pts (sorted by d*): the median of the window groups' widths.
    auto shell_width = [&](size_t lo, size_t hi) -> std::optional<double> {
        std::vector<Point> shell(pts.begin() + lo, pts.begin() + hi);
        std::sort(shell.begin(), shell.end(), [](const Point &a, const Point &b) { return a.window < b.window; });
        std::vector<double> widths;
        for (size_t g = 0; g < WINDOW_GROUPS; ++g) {
            const size_t glo = g * shell.size() / WINDOW_GROUPS, ghi = (g + 1) * shell.size() / WINDOW_GROUPS;
            double sw = 0.0, svar = 0.0, swde2 = 0.0, sde2 = 0.0;
            for (size_t i = glo; i < ghi; ++i) {
                sw += shell[i].w;
                svar += shell[i].var_w;
                swde2 += shell[i].w * shell[i].de2;
                sde2 += shell[i].de2;
            }
            if (!(sw >= MIN_SHELL_SIGNAL * std::sqrt(svar)) || !(swde2 > 0.0))
                continue;
            if (const auto v = UntruncatedVariance(swde2 / sw, sde2 / static_cast<double>(ghi - glo)))
                widths.push_back(*v);
        }
        if (widths.empty())
            return std::nullopt;
        std::sort(widths.begin(), widths.end());
        const size_t m = widths.size();
        return m % 2 ? widths[m / 2] : 0.5 * (widths[m / 2 - 1] + widths[m / 2]);
    };

    // For the three-term formula (the prediction window): shells of equal count in d*^2.
    std::vector<double> sx, sy;
    for (size_t k = 0; k < N_SHELLS; ++k) {
        const size_t lo = k * pts.size() / N_SHELLS, hi = (k + 1) * pts.size() / N_SHELLS;
        const auto y = shell_width(lo, hi);
        if (!y)
            continue;
        double sumx = 0.0;
        for (size_t i = lo; i < hi; ++i)
            sumx += pts[i].x;
        const double x = sumx / static_cast<double>(hi - lo);
        sx.push_back(x);
        sy.push_back(*y);
        logger.Debug("Stills rocking width at d = {:.2f} A: {:.2e} 1/A ({:.3f} deg)",
                     1.0 / std::sqrt(x), std::sqrt(*y), std::sqrt(*y / x) * kRadToDeg);
    }

    // For the partiality: shells of equal width in d*, kept as a table. Equal-count shells crowd where
    // the reflections are - on a run predicting ~1000 reflections a frame, all eight sat at 3.8-1.8 A, so
    // the low-resolution width was the formula's extrapolation, and the three-term fit flipped between
    // forms with the number of frames (to a pure bandwidth at 50k frames). Measured where it is used and
    // interpolated, the width needs no form. A wider curve at finer resolution is the physics (every term
    // grows with d*), so the table is made non-decreasing (pool adjacent violators, weighted by the
    // shells' summed signal).
    std::vector<double> table_dstar, table_s, table_weight;
    {
        const double lo_d = std::sqrt(pts.front().x), hi_d = std::sqrt(pts.back().x);
        size_t lo = 0;
        for (size_t k = 0; k < N_SHELLS; ++k) {
            const double edge = lo_d + (hi_d - lo_d) * static_cast<double>(k + 1) / N_SHELLS;
            size_t hi = lo;
            while (hi < pts.size() && (k + 1 == N_SHELLS || std::sqrt(pts[hi].x) < edge))
                ++hi;
            double sw = 0.0, svar = 0.0, sds = 0.0;
            for (size_t i = lo; i < hi; ++i) {
                sw += pts[i].w;
                svar += pts[i].var_w;
                sds += std::sqrt(pts[i].x);
            }
            if (hi - lo >= 400 && sw >= MIN_SHELL_SIGNAL * std::sqrt(svar))
                if (const auto y = shell_width(lo, hi)) {
                    table_dstar.push_back(sds / static_cast<double>(hi - lo));
                    table_s.push_back(std::sqrt(*y));
                    table_weight.push_back(sw / std::sqrt(svar));
                }
            lo = hi;
        }
        std::vector<double> block_s, block_w;
        std::vector<size_t> block_n;
        for (size_t i = 0; i < table_s.size(); ++i) {
            block_s.push_back(table_s[i]);
            block_w.push_back(table_weight[i]);
            block_n.push_back(1);
            while (block_s.size() > 1 && block_s[block_s.size() - 2] > block_s.back()) {
                const size_t j = block_s.size() - 2;
                block_s[j] = (block_s[j] * block_w[j] + block_s.back() * block_w.back()) / (block_w[j] + block_w.back());
                block_w[j] += block_w.back();
                block_n[j] += block_n.back();
                block_s.pop_back();
                block_w.pop_back();
                block_n.pop_back();
            }
        }
        table_s.clear();
        for (size_t j = 0; j < block_s.size(); ++j)
            table_s.insert(table_s.end(), block_n[j], block_s[j]);
        if (table_dstar.size() < 3) {
            table_dstar.clear();
            table_s.clear();
        }
    }

    if (sx.size() < 3)
        return std::nullopt;

    // Least squares for s^2 = A + B d*^2 + C d*^4 over the shells. Every term is a variance, so a negative
    // one is unphysical: of the fits on each subset of the three terms, the best with none negative is kept.
    std::array<double, 3> best{};
    double best_residual = INFINITY;
    for (int mask = 1; mask < 8; ++mask) {
        std::vector<int> terms;
        for (int t = 0; t < 3; ++t)
            if (mask & (1 << t))
                terms.push_back(t);
        const int m = static_cast<int>(terms.size());
        Eigen::MatrixXd M = Eigen::MatrixXd::Zero(m, m);
        Eigen::VectorXd v = Eigen::VectorXd::Zero(m);
        for (size_t k = 0; k < sx.size(); ++k)
            for (int i = 0; i < m; ++i) {
                const double fi = std::pow(sx[k], terms[i]);
                v(i) += fi * sy[k];
                for (int j = 0; j < m; ++j)
                    M(i, j) += fi * std::pow(sx[k], terms[j]);
            }
        const Eigen::VectorXd c = M.ldlt().solve(v);
        if (!c.allFinite() || (c.array() < 0.0).any())
            continue;
        std::array<double, 3> coef{};
        for (int i = 0; i < m; ++i)
            coef[terms[i]] = c(i);
        double residual = 0.0;
        for (size_t k = 0; k < sx.size(); ++k) {
            const double r = sy[k] - coef[0] - coef[1] * sx[k] - coef[2] * sx[k] * sx[k];
            residual += r * r;
        }
        if (residual < best_residual) {
            best_residual = residual;
            best = coef;
        }
    }
    if (!std::isfinite(best_residual) || !(best[0] + best[1] + best[2] > 0.0))
        return std::nullopt;
    // The d*^4 term is a bandwidth: its radial smear bw |q_z|, with |q_z| = lambda d*^2 / 2 on the sphere.
    const RunWidth width{.gamma0 = std::sqrt(best[0]), .gamma_e = std::sqrt(best[1]),
                         .bandwidth_sigma = 2.0 * std::sqrt(best[2]) / experiment_.GetWavelength_A(),
                         .radial_scale = radial_scale, .table_dstar = table_dstar, .table_s = table_s};
    logger.Info("Stills rocking width measured: {:.2e} 1/A + {:.3f} deg angular + {:.4f} bandwidth (FWHM)",
                width.gamma0, width.gamma_e * kRadToDeg, width.bandwidth_sigma * 2.3548);
    std::string table;
    for (size_t i = 0; i < table_dstar.size(); ++i)
        table += fmt::format(" {:.2f} A {:.2e}", 1.0 / table_dstar[i], table_s[i]);
    logger.Info("Stills rocking width for the partiality, against resolution (1/A):{}",
                table.empty() ? " none, the formula is used" : table);
    return width;
}

double StillsPartialityRefine::Run(std::vector<IntegrationOutcome> &outcomes, size_t nthreads,
                                   const std::optional<UnitCell> &reference_cell,
                                   const std::optional<RunWidth> &given_width,
                                   std::optional<RunWidth> *used_width) const {
    const auto &indexing = experiment_.GetIndexingSettings();
    auto merged_crystal = [&](const IntegrationOutcome &o) {
        return !reference_cell || o.latt.GetUnitCell().is_close(*reference_cell,
                                                                indexing.GetUnitCellDistTolerance(),
                                                                indexing.GetUnitCellAngleTolerance_deg());
    };
    if (nthreads == 0)
        nthreads = AvailableCpus();
    nthreads = std::max<size_t>(1, nthreads);

    double last_mean_tilt = 0.0;
    // Where the frames were predicted through a curve measured before the image pass, the partialities use
    // that same curve. Measured again here, through windows that follow it, the width reads narrower at low
    // resolution (its constant term falls to about zero), and the merge it gave moved R-free against the
    // models by up to +-0.015 per run in both directions, where the one curve moved it by at most +0.008.
    std::optional<RunWidth> run_width = given_width;
    if (!run_width) {
        double radial_scale = 1.0;
        for (int i = 0; i < 3; ++i)
            radial_scale = MeasureRadialScale(outcomes, reference_cell, radial_scale);
        Logger("StillsPartialityRefine").Info("Stills rocking curve centre: reciprocal lattice scaled by {:+.2e}",
                                              radial_scale - 1.0);
        run_width = MeasureRunWidth(outcomes, reference_cell, radial_scale);
    }
    const double radial_scale = run_width ? run_width->radial_scale : 1.0;
    if (used_width)
        *used_width = run_width;

    for (int outer = 0; outer < settings_.outer_iterations; ++outer) {
        // Reference full intensities from the current corrections.
        const std::vector<MergedReflection> merged = MergeAll(experiment_, outcomes, nthreads, reference_cell);
        std::map<HKLKey, double> reference;
        for (const auto &m: merged)
            reference[hkl_key_generator_(m)] = m.I;

        std::atomic<double> tilt_sum{0.0};
        std::atomic<size_t> tilt_n{0};
        std::atomic<size_t> next{0};

        auto worker = [&]() {
            size_t i = next.fetch_add(1);
            while (i < outcomes.size()) {
                const double t = merged_crystal(outcomes[i]) ? RefineOne(outcomes[i], reference, run_width, radial_scale) : 0.0;
                if (t > 0.0) {
                    double prev = tilt_sum.load();
                    while (!tilt_sum.compare_exchange_weak(prev, prev + t)) {}
                    tilt_n.fetch_add(1);
                }
                i = next.fetch_add(1);
            }
        };

        const size_t nt = std::min(nthreads, std::max<size_t>(1, outcomes.size()));
        if (nt <= 1) {
            worker();
        } else {
            std::vector<std::future<void>> futures;
            futures.reserve(nt);
            for (size_t t = 0; t < nt; ++t)
                futures.emplace_back(std::async(std::launch::async, worker));
            for (auto &f: futures)
                f.get();
        }

        last_mean_tilt = tilt_n > 0 ? tilt_sum.load() / static_cast<double>(tilt_n.load()) : 0.0;
    }

    return last_mean_tilt;
}