Jungfraujoch/image_analysis/scale_merge/Merge.cpp

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

#include "Merge.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <random>
#include <unordered_map>

#include <spdlog/fmt/fmt.h>

#include <gemmi/reciproc.hpp>

#include "../../common/CorrelationCoefficient.h"
#include "../../common/ResolutionShells.h"
#include "../../common/Definitions.h"
#include "HKLKey.h"
#include "RfreeFlags.h"
#include "FrenchWilson.h"

namespace {
    // Deterministic CC1/2 half-set assignment: a splitmix64 bit-mix of the image's stable index.
    // A pure function of image identity (not a draw from a shared RNG in call order) keeps the split
    // reproducible run-to-run, independent of AddImage call order, and safe under concurrent merging.
    int HalfForImage(int64_t image_id) {
        uint64_t z = static_cast<uint64_t>(image_id) + 0x9e3779b97f4a7c15ULL;
        z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
        z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
        z = z ^ (z >> 31);
        return static_cast<int>(z & 1ULL);
    }
}

MergeOnTheFly::MergeOnTheFly(const DiffractionExperiment &x)
    : space_group_number(x.GetSpaceGroupNumber().value_or(1)),
      scaling_settings(x.GetScalingSettings()),
      indexing_settings(x.GetIndexingSettings()),
      high_resolution_limit(scaling_settings.GetHighResolutionLimit_A()),
      // A min-image-CC of 0 (the default) means "no limit": leave the optional
      // empty so the per-image CC cut is inactive. Otherwise a 0.0 threshold
      // would silently drop every image with a non-positive per-image CC (which
      // also wrongly zeroed N_obs in MergeStats, since it masks with cc_mask=true
      // while the merge keeps all images).
      image_cc_limit(scaling_settings.GetMinCCForImage() > 0.0
                         ? std::optional<double>(scaling_settings.GetMinCCForImage())
                         : std::nullopt),
      min_partiality(scaling_settings.GetMinPartiality()),
      generator(scaling_settings.GetMergeFriedel(), space_group_number),
      reject_outliers(scaling_settings.GetOutlierRejectNsigma() > 0.0),
      reject_nsigma(scaling_settings.GetOutlierRejectNsigma()) {
}

MergeOnTheFly &MergeOnTheFly::ReferenceCell(const std::optional<UnitCell> &cell) {
    reference_cell = cell;
    return *this;
}

bool MergeOnTheFly::IsMaskedRing(const Reflection &r) const {
    if (masked_ice_rings.empty())
        return false;
    const int ring = IceRingIndex(r.d, mask_ice_half_width_q);
    return ring >= 0 && ring < static_cast<int>(masked_ice_rings.size()) && masked_ice_rings[ring];
}

void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id, bool cc_mask) {
    std::unique_lock ul(merged_mutex);

    if (Mask(outcome, cc_mask))
        return;

    const int half = HalfForImage(image_id);

    for (const auto &r: outcome.reflections) {
        if (generator.IsSystematicallyAbsent(r))
            continue;

        if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr))
            continue;
        if (!AcceptReflection(r, high_resolution_limit))
            continue;
        if (exclude_ice_rings && r.on_ice_ring)
            continue;
        if (IsMaskedRing(r))
            continue;
        if (r.partiality < min_partiality)
            continue;

        const float I_corr = r.I * r.image_scale_corr;
        float sigma_corr = r.sigma * r.image_scale_corr;
        if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0)
            continue;
        auto hkl = generator(r);
        auto hkl_key = hkl.pack();
        sigma_corr = CorrectedSigma(I_corr, sigma_corr, hkl_key, r.partiality);

        // Robust outlier rejection: drop this observation if it sits more than
        // reject_nsigma error-model sigmas from the reflection's median. Needs the active
        // error model so sigma_corr reflects the real scatter (else the threshold is the
        // bare counting sigma and would cull good partials).
        if (reject_outliers && error_model_active) {
            const auto mit = reject_median_I.find(hkl_key);
            if (mit != reject_median_I.end() &&
                std::fabs(I_corr - mit->second) > reject_nsigma * sigma_corr) {
                ++reject_count;
                continue;
            }
        }

        auto it = accumulator.find(hkl_key);
        if (it == accumulator.end())
            it = accumulator.emplace(hkl_key, MergeAccum{
                                 .h = hkl.plus ? hkl.h : -hkl.h,
                                 .k = hkl.plus ? hkl.k : -hkl.k,
                                 .l = hkl.plus ? hkl.l : -hkl.l,
                             }).first;

        const float w = 1.0f / (sigma_corr * sigma_corr);
        const float wI = w * I_corr;

        it->second.sum_wI += wI;
        it->second.sum_w += w;

        it->second.sum_wI_half[half] += wI;
        it->second.sum_w_half[half] += w;
        it->second.n_half[half]++;

        if (!std::isfinite(it->second.d) && std::isfinite(r.d) && r.d > 0.0f)
            it->second.d = r.d;
    }
}

double MergeOnTheFly::RefineModulation(std::vector<IntegrationOutcome> &outcomes) {
    // A minimum held-out generalizing gain (fraction of the held-out scatter) before the surface is
    // applied - a margin, so a noise-level "improvement" never engages the correction. Matches the
    // rotation ApplyCellSurface gate.
    constexpr double CV_MIN_RELATIVE_GAIN = 0.02;
    constexpr int NB = 16;
    const int ncell = NB * NB;

    // One scaled observation reduced to what the surface fit needs. `parity` (image index & 1) drives the
    // even/odd cross-validation split - the stills analogue of the rotation frame parity.
    struct MObs { double I, sigma, corr; float px, py; int32_t group; int parity; int cell; };
    std::vector<MObs> obs;

    // Accept exactly what AddImage merges (systematic absence, scale/resolution/ice/partiality filters),
    // and additionally require a finite detector position. Returns the dense ASU-group id or -1.
    std::unordered_map<uint64_t, int> group_of;
    auto accept = [&](const Reflection &r, MObs &out) -> bool {
        if (generator.IsSystematicallyAbsent(r)) return false;
        if (r.image_scale_corr <= 0.0f || !std::isfinite(r.image_scale_corr)) return false;
        if (!AcceptReflection(r, high_resolution_limit)) return false;
        if (exclude_ice_rings && r.on_ice_ring) return false;
        if (IsMaskedRing(r)) return false;
        if (r.partiality < min_partiality) return false;
        if (!std::isfinite(r.predicted_x) || !std::isfinite(r.predicted_y)) return false;
        const float I_corr = r.I * r.image_scale_corr, sigma_corr = r.sigma * r.image_scale_corr;
        if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return false;
        const uint64_t key = generator(r).pack();
        auto [it, inserted] = group_of.try_emplace(key, static_cast<int>(group_of.size()));
        out.I = r.I; out.sigma = r.sigma; out.corr = r.image_scale_corr;
        out.px = r.predicted_x; out.py = r.predicted_y; out.group = it->second;
        return true;
    };

    float pxmin = std::numeric_limits<float>::infinity(), pxmax = -pxmin, pymin = pxmin, pymax = -pxmin;
    for (size_t i = 0; i < outcomes.size(); ++i) {
        const int parity = static_cast<int>(i & 1);
        for (const auto &r : outcomes[i].reflections) {
            MObs m{};
            if (!accept(r, m)) continue;
            m.parity = parity;
            pxmin = std::min(pxmin, m.px); pxmax = std::max(pxmax, m.px);
            pymin = std::min(pymin, m.py); pymax = std::max(pymax, m.py);
            obs.push_back(m);
        }
    }
    const int n_groups = static_cast<int>(group_of.size());
    if (obs.size() < static_cast<size_t>(8 * ncell) || !(pxmax > pxmin) || !(pymax > pymin))
        return 0.0;   // too sparse to over-determine a 16x16 surface, or degenerate detector footprint

    const float sx = NB / (pxmax - pxmin), sy = NB / (pymax - pymin);
    for (auto &m : obs) {
        const int ix = std::clamp(static_cast<int>((m.px - pxmin) * sx), 0, NB - 1);
        const int iy = std::clamp(static_cast<int>((m.py - pymin) * sy), 0, NB - 1);
        m.cell = ix * NB + iy;
    }

    // Fit the per-cell factor over {parity subset} (parity < 0 = all obs), n_iter alternating rounds against
    // that subset's own inverse-variance reference: Tikhonov pull to 1, gauge-fixed to a den-weighted
    // geometric mean of 1 so it never drifts the overall scale. Mirrors RotationScaleMerge::ApplyCellSurface.
    constexpr int N_ITER = 3;
    auto fit_surface = [&](int parity) -> std::vector<double> {
        std::vector<double> A(ncell, 1.0);
        for (int it = 0; it < N_ITER; ++it) {
            std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
            for (const auto &o : obs) {
                if (parity >= 0 && o.parity != parity) continue;
                const double a = A[o.cell], sc = o.sigma * o.corr * a, w = 1.0 / (sc * sc);
                sw[o.group] += w; swI[o.group] += w * o.I * o.corr * a;
            }
            std::vector<double> num(ncell, 0.0), den(ncell, 0.0);
            for (const auto &o : obs) {
                if ((parity >= 0 && o.parity != parity) || sw[o.group] <= 0.0) continue;
                const double Iref = swI[o.group] / sw[o.group], a = A[o.cell];
                const double Is = o.I * o.corr * a, sc = o.sigma * o.corr * a;
                if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0) || !(sc > 0.0)) continue;
                const double w = 1.0 / (sc * sc);
                num[o.cell] += w * Is * Iref; den[o.cell] += w * Is * Is;
            }
            std::vector<double> dsorted = den;
            std::nth_element(dsorted.begin(), dsorted.begin() + dsorted.size() / 2, dsorted.end());
            const double lambda = 0.1 * std::max(1e-30, dsorted[dsorted.size() / 2]);
            double logsum = 0.0, wsum = 0.0;
            std::vector<double> upd(ncell, 1.0);
            for (int c = 0; c < ncell; ++c) upd[c] = (num[c] + lambda) / (den[c] + lambda);
            for (int c = 0; c < ncell; ++c) if (den[c] > 0.0) { logsum += den[c] * std::log(upd[c]); wsum += den[c]; }
            const double gm = wsum > 0.0 ? std::exp(logsum / wsum) : 1.0;
            for (int c = 0; c < ncell; ++c) A[c] = std::clamp(A[c] * upd[c] / gm, 0.25, 4.0);
        }
        return A;
    };
    // Sigma-independent (R-meas-like) agreement of the held-out equivalents: sum|Is - Iref| / sum|Iref|.
    // A fractional metric cannot be gamed by a surface that merely reshapes sigma via corr.
    auto score = [&](int parity, const std::vector<double> &A) -> double {
        std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
        for (const auto &o : obs) {
            if (o.parity != parity) continue;
            const double a = A[o.cell], Is = o.I * o.corr * a, sc = o.sigma * o.corr * a, w = 1.0 / (sc * sc);
            sw[o.group] += w; swI[o.group] += w * Is;
        }
        double num = 0.0, den = 0.0;
        for (const auto &o : obs) {
            if (o.parity != parity || sw[o.group] <= 0.0) continue;
            const double a = A[o.cell], Is = o.I * o.corr * a, Iref = swI[o.group] / sw[o.group];
            if (!std::isfinite(Iref) || Iref <= 0.0) continue;
            num += std::abs(Is - Iref); den += Iref;
        }
        return den > 0.0 ? num / den : 0.0;
    };

    // Cross-validate: fit on even images, score the held-out odd equivalents (and vice versa). Apply the
    // full-data surface only if the held-out agreement improves by a clear margin.
    const std::vector<double> ident(ncell, 1.0);
    const std::vector<double> A_even = fit_surface(0), A_odd = fit_surface(1);
    const double base = score(1, ident) + score(0, ident);
    const double gain = base - (score(1, A_even) + score(0, A_odd));
    if (!(gain > CV_MIN_RELATIVE_GAIN * base))
        return 0.0;   // not cross-validated: the correction stays a no-op (caller logs)
    const std::vector<double> A = fit_surface(-1);

    // Fold the surface into each accepted reflection's image_scale_corr (recompute its cell deterministically).
    for (auto &outcome : outcomes)
        for (auto &r : outcome.reflections) {
            MObs m{};
            if (!accept(r, m)) continue;
            const int ix = std::clamp(static_cast<int>((m.px - pxmin) * sx), 0, NB - 1);
            const int iy = std::clamp(static_cast<int>((m.py - pymin) * sy), 0, NB - 1);
            r.image_scale_corr = static_cast<float>(r.image_scale_corr * A[ix * NB + iy]);
        }
    return gain / std::max(base, 1e-30);   // held-out gain fraction (caller logs)
}

float MergeOnTheFly::CorrectedSigma(float I_corr, float sigma_corr, uint64_t hkl_key, float partiality) const {
    if (!error_model_active)
        return sigma_corr;

    // Intensity for the (b*I)^2 term: the reflection's mean (constant over its
    // observations), falling back to this observation only if the mean is unknown.
    const auto it = error_model_mean_I.find(hkl_key);
    const double I_for_b = (it != error_model_mean_I.end()) ? it->second : I_corr;

    double v = error_model_a * static_cast<double>(sigma_corr) * sigma_corr
               + (error_model_b * I_for_b) * (error_model_b * I_for_b);

    // Partiality-model uncertainty: a reflection recorded at fraction p carries a systematic intensity
    // error ~ (dp/p) that is proportional to <I> and grows as p falls - plain counting sigma misses it,
    // so strong low-p partials would otherwise be over-trusted. This is the stills-partiality analog of
    // the rotation --capture-uncertainty term ((1-captured_fraction)*I in RotationScaleMerge). Inert when
    // partiality == 1 (no stills partiality model). Gated on a real systematic (error_model_b > 1, i.e.
    // ISa < 1): on weak counting-limited data (small b) it would only over-concentrate the merge and hurt.
    const double c = scaling_settings.GetPartialityUncertaintyCoeff();
    if (c > 0.0 && error_model_b > 1.0) {
        const double one_minus_p = std::max(0.0, std::min(1.0, 1.0 - static_cast<double>(partiality)));
        const double t = c * I_for_b * one_minus_p;
        v += t * t;
    }
    return (v > 0.0) ? static_cast<float>(std::sqrt(v)) : sigma_corr;
}

void MergeOnTheFly::RefineErrorModel(const std::vector<IntegrationOutcome> &outcomes) {
    // Reset to identity up front: every early return below then leaves the model
    // inactive (CorrectedSigma returns sigma unchanged) rather than keeping a stale
    // a/b from a previous call alongside a freshly-cleared mean map.
    // Median of the chi-square(1) distribution: a single observation's squared deviation from its
    // reflection mean, divided by its variance, is chi-square(1)-distributed, so its median is this
    // fraction of its mean. Used both to de-bias the median-based variance fit and to normalize the
    // reported median reduced chi^2 so that honestly calibrated sigmas give 1.0 (not 0.4549).
    constexpr double CHI2_1_MEDIAN = 0.454936;

    error_model_active = false;
    error_model_a = 1.0;
    error_model_b = 0.0;
    error_model_chi2 = 0.0;
    error_model_mean_I.clear();
    reject_median_I.clear();
    reject_count = 0;

    // --- 1. Collect accepted, scaled observations grouped by symmetry-equivalent hkl,
    //        applying exactly the filters AddImage uses. ---
    struct Obs { float I, sigma; };
    std::unordered_map<uint64_t, std::vector<Obs>> groups;

    for (const auto &outcome: outcomes) {
        if (Mask(outcome, false))
            continue;
        for (const auto &r: outcome.reflections) {
            if (generator.IsSystematicallyAbsent(r))
                continue;
            if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr))
                continue;
            if (!AcceptReflection(r, high_resolution_limit))
                continue;
            if (exclude_ice_rings && r.on_ice_ring)
                continue;
            if (IsMaskedRing(r))
                continue;
            if (r.partiality < min_partiality)
                continue;
            const float I_corr = r.I * r.image_scale_corr;
            const float sigma_corr = r.sigma * r.image_scale_corr;
            if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f)
                continue;
            groups[generator(r).pack()].push_back({I_corr, sigma_corr});
        }
    }

    // --- 2. One global pool of (sigma^2, <I>^2, bias-corrected squared deviation). For an
    //        observation in a group of n, the residual from the inverse-variance mean has
    //        E[(I_i - <I>)^2] = sigma_i^2 (1 - h_i), h_i = w_i / sum_w (its leverage). The
    //        (b*I)^2 term uses the reflection mean, so the mean (not I_i) is the abscissa. ---
    struct Sample { double s2, I2, dev2; };
    std::vector<Sample> samples;

    for (const auto &[key, obs]: groups) {
        if (obs.size() < 2)
            continue;
        double sum_w = 0.0, sum_wI = 0.0;
        for (const auto &o: obs) {
            const double w = 1.0 / (static_cast<double>(o.sigma) * o.sigma);
            sum_w += w;
            sum_wI += w * o.I;
        }
        if (!(sum_w > 0.0))
            continue;
        const double mean = sum_wI / sum_w;
        error_model_mean_I[key] = static_cast<float>(mean);
        // Robust centre for outlier rejection: the median intensity (resists the very
        // outliers the inverse-variance mean is being protected from). Only when active.
        if (reject_outliers) {
            std::vector<float> iv;
            iv.reserve(obs.size());
            for (const auto &o: obs)
                iv.push_back(o.I);
            std::nth_element(iv.begin(), iv.begin() + iv.size() / 2, iv.end());
            reject_median_I[key] = iv[iv.size() / 2];
        }
        const double I2 = mean * mean;
        for (const auto &o: obs) {
            const double w = 1.0 / (static_cast<double>(o.sigma) * o.sigma);
            const double factor = 1.0 - w / sum_w;
            if (factor < 0.05)
                continue;
            const double resid = static_cast<double>(o.I) - mean;
            samples.push_back({static_cast<double>(o.sigma) * o.sigma, I2, resid * resid / factor});
        }
    }

    // --- 3. Fit global dev2 = a*sigma^2 + b^2*<I>^2. Bin by intensity (the per-observation
    //        dev2 is chi-square-1 noisy) and take medians; weight the bins by 1/dev2^2 so it
    //        is a *relative* fit - otherwise the strong bins (which fix b) swamp the weak
    //        bins (which fix a) and the weak sigmas stay over-confident. ---
    constexpr int n_bins = 16;
    if (samples.size() < static_cast<size_t>(8 * n_bins))
        return; // too little multiplicity to fit -> leave identity

    std::sort(samples.begin(), samples.end(),
              [](const Sample &p, const Sample &q) { return p.I2 < q.I2; });

    auto median = [](std::vector<double> &v) {
        std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
        return v[v.size() / 2];
    };

    // Per-intensity-bin medians of (sigma^2, <I>^2, dev2).
    std::vector<double> bs2, bI2, bd2;
    bs2.reserve(n_bins); bI2.reserve(n_bins); bd2.reserve(n_bins);
    const size_t per = samples.size() / n_bins;
    for (int bin = 0; bin < n_bins; ++bin) {
        const size_t lo = bin * per;
        const size_t hi = (bin == n_bins - 1) ? samples.size() : lo + per;
        std::vector<double> vs2, vI2, vd2;
        vs2.reserve(hi - lo); vI2.reserve(hi - lo); vd2.reserve(hi - lo);
        for (size_t i = lo; i < hi; ++i) {
            vs2.push_back(samples[i].s2);
            vI2.push_back(samples[i].I2);
            vd2.push_back(samples[i].dev2);
        }
        bs2.push_back(median(vs2));
        bI2.push_back(median(vI2));
        // The per-observation dev2 is sigma^2 * chi-square(1)-distributed, whose MEDIAN is 0.4549 of
        // its mean. Fitting the model to the robust median would therefore calibrate the variances to
        // 0.4549x their true value (reduced chi^2 ~ 1/0.4549 = 2.2). Divide the median by that constant
        // to recover an unbiased estimate of the mean (E[dev2] = sigma^2), keeping the robustness of
        // the median while targeting reduced chi^2 = 1.
        bd2.push_back(median(vd2) / CHI2_1_MEDIAN);
    }

    // Relative-weighted (1/dev2^2) least squares for (a, b^2). Floor the weight's dev2 at a
    // small fraction of the typical bin dev2: an absolute floor (1e-30) does not stop a
    // near-zero-scatter bin from acquiring a runaway weight and hijacking the fit, so the
    // floor must scale with the data. The regression target keeps the unfloored dev2.
    std::vector<double> bd2_sorted = bd2;
    const double dev2_floor = std::max(1e-30, 1e-3 * median(bd2_sorted));
    double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0;
    for (int bin = 0; bin < n_bins; ++bin) {
        const double s2 = bs2[bin], I2 = bI2[bin], d2 = bd2[bin];
        const double d2w = std::max(d2, dev2_floor);
        const double wgt = 1.0 / (d2w * d2w);
        Ass += wgt * s2 * s2;
        AsI += wgt * s2 * I2;
        AII += wgt * I2 * I2;
        Bs  += wgt * s2 * d2;
        BI  += wgt * I2 * d2;
    }

    // Reject a near-collinear (ill-conditioned) system *relatively*: det lies in
    // [0, Ass*AII] by Cauchy-Schwarz, so compare against that scale rather than 1e-30.
    const double det = Ass * AII - AsI * AsI;
    if (!(det > 1e-10 * Ass * AII))
        return;
    const double a  = std::clamp((Bs * AII - BI * AsI) / det, 0.25, 100.0);
    const double b2 = std::max((Ass * BI - AsI * Bs) / det, 0.0);

    error_model_a = a;
    error_model_b = std::sqrt(b2);
    error_model_active = true;

    // Achieved goodness of fit: the median of the per-observation dev2/(a*sigma^2 + (b*<I>)^2). That
    // ratio is chi-square(1)-distributed (median 0.4549) when the sigmas are correct, so normalize by
    // CHI2_1_MEDIAN to report a median reduced chi^2 that targets 1.0. The median (not mean) keeps it
    // robust to the heavy outlier tail of serial data.
    std::vector<double> chi2;
    chi2.reserve(samples.size());
    for (const auto &s: samples) {
        const double v = a * s.s2 + b2 * s.I2;
        if (v > 0.0)
            chi2.push_back(s.dev2 / v);
    }
    error_model_chi2 = chi2.empty() ? 0.0 : median(chi2) / CHI2_1_MEDIAN;
}

// Per-crystal CC1/2-delta rejection (CrystFEL deltaCChalf style). Each image is assigned
// to one CC1/2 half, so removing an image only perturbs that half's per-reflection means.
// deltaCChalf_i = CC1/2(all) - CC1/2(without image i): a NEGATIVE value means removing the
// image RAISES CC1/2, i.e. it is inconsistent with the consensus. We flag images whose
// deltaCChalf is a low-side statistical outlier (< mean - nsigma*stddev). Reference-free.
// Two passes over the (retained) outcomes; per-image contributions are re-derived, not
// stored, so memory stays O(unique reflections + images) for full 200k-frame datasets.
std::vector<char> MergeOnTheFly::DeltaCChalfReject(const std::vector<IntegrationOutcome> &outcomes,
                                                   double nsigma) const {
    struct Acc { double swI[2] = {0, 0}; double sw[2] = {0, 0}; size_t n[2] = {0, 0}; };
    std::unordered_map<uint64_t, Acc> acc;
    std::vector<int> img_half(outcomes.size(), 0);

    // ---- pass 1: accumulate half-set sums, record each image's half ----
    auto contribution = [&](const Reflection &r, uint64_t &key, double &wI, double &w) -> bool {
        if (generator.IsSystematicallyAbsent(r)) return false;
        if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) return false;
        if (!AcceptReflection(r, high_resolution_limit)) return false;
        if (r.partiality < min_partiality) return false;
        const double I = static_cast<double>(r.I) * r.image_scale_corr;
        const double s = static_cast<double>(r.sigma) * r.image_scale_corr;
        if (!std::isfinite(I) || !std::isfinite(s) || s <= 0.0) return false;
        w = 1.0 / (s * s);
        wI = w * I;
        key = generator(r).pack();
        return true;
    };

    for (size_t i = 0; i < outcomes.size(); ++i) {
        // Same deterministic half-set as the merge (HalfForImage), so deltaCChalf is measured on the
        // exact CC1/2 split the statistics report - not an independent, order-dependent RNG draw.
        const int half = HalfForImage(static_cast<int64_t>(i));
        img_half[i] = half;
        for (const auto &r : outcomes[i].reflections) {
            uint64_t key; double wI, w;
            if (!contribution(r, key, wI, w)) continue;
            auto &a = acc[key];
            a.swI[half] += wI; a.sw[half] += w; a.n[half]++;
        }
    }

    // ---- baseline Pearson over reflections present in BOTH halves (x=half0, y=half1) ----
    auto pearson = [](double N, double Sx, double Sy, double Sxx, double Syy, double Sxy) -> double {
        const double cov = N * Sxy - Sx * Sy;
        const double vx = N * Sxx - Sx * Sx, vy = N * Syy - Sy * Sy;
        const double den = std::sqrt(vx * vy);
        return den > 0.0 ? cov / den : 0.0;
    };
    double N = 0, Sx = 0, Sy = 0, Sxx = 0, Syy = 0, Sxy = 0;
    for (const auto &kv : acc) {
        const auto &a = kv.second;
        if (a.n[0] == 0 || a.n[1] == 0) continue;
        const double x = a.swI[0] / a.sw[0], y = a.swI[1] / a.sw[1];
        N += 1; Sx += x; Sy += y; Sxx += x * x; Syy += y * y; Sxy += x * y;
    }
    const double cc_base = pearson(N, Sx, Sy, Sxx, Syy, Sxy);

    // ---- pass 2: leave-one-out deltaCChalf per image ----
    std::vector<double> delta(outcomes.size(), 0.0);
    for (size_t i = 0; i < outcomes.size(); ++i) {
        const int h = img_half[i];
        // aggregate this image's contributions per reflection key (an image may, rarely,
        // touch the same ASU reflection twice)
        std::unordered_map<uint64_t, std::pair<double, double>> mine; // key -> (sum wI, sum w), count via .first
        std::unordered_map<uint64_t, size_t> mine_n;
        for (const auto &r : outcomes[i].reflections) {
            uint64_t key; double wI, w;
            if (!contribution(r, key, wI, w)) continue;
            auto &p = mine[key]; p.first += wI; p.second += w; mine_n[key]++;
        }
        double n = N, sx = Sx, sy = Sy, sxx = Sxx, syy = Syy, sxy = Sxy;
        for (const auto &m : mine) {
            const auto &a = acc.at(m.first);
            if (a.n[0] == 0 || a.n[1] == 0) continue; // reflection not in CC1/2
            const double x0 = a.swI[0] / a.sw[0], y0 = a.swI[1] / a.sw[1];
            n -= 1; sx -= x0; sy -= y0; sxx -= x0 * x0; syy -= y0 * y0; sxy -= x0 * y0;
            const double swI_h = a.swI[h] - m.second.first;
            const double sw_h = a.sw[h] - m.second.second;
            if (a.n[h] - mine_n[m.first] == 0 || sw_h <= 0.0) continue; // reflection drops half h
            const double mean_h = swI_h / sw_h;
            const double xnew = (h == 0) ? mean_h : x0;
            const double ynew = (h == 1) ? mean_h : y0;
            n += 1; sx += xnew; sy += ynew; sxx += xnew * xnew; syy += ynew * ynew; sxy += xnew * ynew;
        }
        delta[i] = cc_base - pearson(n, sx, sy, sxx, syy, sxy);
    }

    // ---- reject low-side outliers: delta < mean - nsigma*stddev ----
    double dm = 0, dv = 0;
    for (double d : delta) dm += d;
    dm /= std::max<size_t>(1, delta.size());
    for (double d : delta) dv += (d - dm) * (d - dm);
    const double dstd = std::sqrt(dv / std::max<size_t>(1, delta.size()));
    const double cut = dm - nsigma * dstd;

    std::vector<char> reject(outcomes.size(), 0);
    for (size_t i = 0; i < outcomes.size(); ++i)
        reject[i] = (outcomes[i].reflections.empty() ? 0 : (delta[i] < cut ? 1 : 0));
    return reject;
}

bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome, bool cc_mask) {
    if (reference_cell) {
        auto cell = outcome.latt.GetUnitCell();
        if (!cell.is_close(*reference_cell,
                          indexing_settings.GetUnitCellDistTolerance(),
                          indexing_settings.GetUnitCellAngleTolerance_deg()))
            return true;
    }

    if (cc_mask && image_cc_limit) {
        if (!outcome.image_scale_cc
            || std::isnan(outcome.image_scale_cc.value())
            || outcome.image_scale_cc.value() < image_cc_limit.value())
            return true;
    }
    return false;
}

std::vector<MergedReflection> MergeOnTheFly::ExportReflections() {
    std::unique_lock ul(merged_mutex);

    std::vector<MergedReflection> out;
    out.reserve(accumulator.size());
    for (const auto &accum: accumulator | std::views::values) {
        if (accum.sum_w <= 0.0)
            continue;

        MergedReflection mr{
            .h = accum.h,
            .k = accum.k,
            .l = accum.l,
            .I = static_cast<float>(accum.sum_wI / accum.sum_w),
            .sigma = SigmaWithSystematicFloor(1.0 / std::sqrt(accum.sum_w),
                                              static_cast<float>(accum.sum_wI / accum.sum_w), error_model_b),
            .I_half = {NAN, NAN},
            .sigma_half = {NAN, NAN},
            .d = accum.d
        };

        if (accum.n_half[0] + accum.n_half[1] > 0 && accum.sum_w_half[0] > 0.0 && accum.sum_w_half[1] > 0.0) {
            for (int i = 0; i < 2; ++i) {
                mr.I_half[i] = static_cast<float>(accum.sum_wI_half[i] / accum.sum_w_half[i]);
                mr.sigma_half[i] = SigmaWithSystematicFloor(1.0 / std::sqrt(accum.sum_w_half[i]),
                                                            mr.I_half[i], error_model_b);
            }
        }

        if (!std::isfinite(accum.d) || accum.d <= 0.0f)
            continue;

        out.emplace_back(mr);
    }

    AssignRfreeFlags(out, space_group_number, scaling_settings.GetRfreeFraction());
    ApplyFrenchWilson(out, space_group_number);
    return out;
}

std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
                                       const std::vector<IntegrationOutcome> &integration_outcome,
                                       bool mask) {
    MergeOnTheFly merge(x);
    for (size_t i = 0; i < integration_outcome.size(); ++i)
        merge.AddImage(integration_outcome[i], static_cast<int64_t>(i), mask);
    return merge.ExportReflections();
}

struct ShellAccum {
    int total_obs = 0;
    int unique = 0;
    int possible = 0;

    double sum_i_over_sigma = 0.0;
    int n_i_over_sigma = 0;

    CorrelationCoefficient cc_half;
    CorrelationCoefficient cc_ref;
};

void CalcPossibleReflections(int space_group_number ,
                             const UnitCell &cell,
                             double d_min,
                             double d_max,
                             const ResolutionShells &shells,
                             std::vector<ShellAccum> &acc,
                             bool merge_friedel) {
    gemmi::UnitCell gemmi_cell = cell;
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
    if (sg == nullptr)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
            "Invalid space group number " + std::to_string(space_group_number));

    // Generate unique reflections
    std::vector<gemmi::Miller> possible_hkls = gemmi::make_miller_vector(gemmi_cell, sg, d_min, d_max, true);
    const gemmi::GroupOps gops = sg->operations();
    CrystalLattice lattice(cell);
    const auto astar = lattice.Astar();
    const auto bstar = lattice.Bstar();
    const auto cstar = lattice.Cstar();

    for (const auto &hkl: possible_hkls) {
        const auto q = hkl[0] * astar + hkl[1] * bstar + hkl[2] * cstar;
        const auto qlen = q.Length();
        if (qlen < 1e-6)
            continue;
        const auto d = 1.0 / qlen;
        const auto shell = shells.GetShell(d);
        if (!shell.has_value())
            continue;
        const int s = *shell;
        if (s >= 0 && s < acc.size())
            // Anomalous (no Friedel merge): an acentric reflection has two unique members (I+ and I-),
            // a centric one only one — match how unique_reflections is counted, so completeness stays
            // <=100% instead of approaching 200%.
            acc[s].possible += (merge_friedel || gops.is_reflection_centric(hkl)) ? 1 : 2;
    }
}


MergeStatistics MergeOnTheFly::MergeStats(const std::vector<MergedReflection> &merged,
                           const std::vector<IntegrationOutcome > &integration_outcome,
                           const std::vector<MergedReflection> &reference,
                           std::optional<double> d_min_override) {

    const int n_shells = scaling_settings.GetReportShellCount();

    auto d_min_limit_A = d_min_override.has_value()
        ? d_min_override : scaling_settings.GetHighResolutionLimit_A();

    std::unordered_map<uint64_t, float> reference_intensities;
    if (!reference.empty()) {
        reference_intensities.reserve(reference.size());
        for (const auto &r: reference) {
            if (!std::isfinite(r.I))
                continue;

            const auto hkl = generator(r);
            reference_intensities[hkl.pack()] = r.I;
        }
    }

    float d_min = std::numeric_limits<float>::max();
    float d_max = 0.0f;

    for (const auto &m: merged) {
        if (!std::isfinite(m.d) || m.d <= 0.0f)
            continue;
        if (d_min_limit_A && m.d < d_min_limit_A)
            continue;

        d_min = std::min(d_min, m.d);
        d_max = std::max(d_max, m.d);
    }

    if (!(d_min < d_max && d_min > 0.0f))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
            "MergeStats: Error in resolution calculation");

    const float d_min_pad = d_min * 0.999f;
    const float d_max_pad = d_max * 1.001f;

    ResolutionShells shells(d_min_pad, d_max_pad, n_shells);
    const auto shell_mean_1_d2 = shells.GetShellMeanOneOverResSq();
    const auto shell_min_res = shells.GetShellMinRes();

    std::vector<ShellAccum> acc(n_shells);

    if (reference_cell.has_value())
        CalcPossibleReflections(space_group_number, reference_cell.value(),
            d_min_pad, d_max_pad, shells, acc, scaling_settings.GetMergeFriedel());

    CorrelationCoefficient cc_half_overall;
    CorrelationCoefficient cc_ref_overall;

    for (const auto &m: merged) {
        const auto shell = shells.GetShell(m.d);
        if (!shell.has_value())
            continue;

        const int s = *shell;
        if (s >= 0 && s < n_shells) {
            if (std::isfinite(m.I) && std::isfinite(m.sigma) && m.sigma > 0.0) {
                acc[s].unique++;
                acc[s].sum_i_over_sigma += m.I / m.sigma;
                ++acc[s].n_i_over_sigma;

                if (!reference_intensities.empty()) {
                    const auto hkl = generator(m);
                    const auto ref_it = reference_intensities.find(hkl.pack());
                    if (ref_it != reference_intensities.end() && std::isfinite(ref_it->second)) {
                        acc[s].cc_ref.Add(m.I, ref_it->second);
                        cc_ref_overall.Add(m.I, ref_it->second);
                    }
                }

                if (std::isfinite(m.I_half[0]) && std::isfinite(m.I_half[1])) {
                    acc[s].cc_half.Add(m.I_half[0], m.I_half[1]);
                    cc_half_overall.Add(m.I_half[0], m.I_half[1]);
                }

            }
        }
    }

    // Per-reflection mean <I>, and a per-reflection accumulator for R_meas - it needs |I_i - <I>|,
    // so the observations are visited again now that the means are known.
    std::unordered_map<uint64_t, float> merged_I;
    merged_I.reserve(merged.size());
    for (const auto &m: merged)
        if (std::isfinite(m.I))
            merged_I[generator(m).pack()] = m.I;

    struct RmeasObs { double sum_abs_dev = 0.0; double sum_I = 0.0; int n = 0; int shell = -1; };
    std::unordered_map<uint64_t, RmeasObs> rmeas_obs;
    rmeas_obs.reserve(merged.size());

    for (int i = 0; i < integration_outcome.size(); ++i) {
        if (Mask(integration_outcome[i], true))
            continue;

        for (const auto &r: integration_outcome[i].reflections) {
            if (generator.IsSystematicallyAbsent(r))
                continue;
            if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr))
                continue;
            if (!AcceptReflection(r, d_min_limit_A))
                continue;
            if (r.partiality < min_partiality)
                continue;

            const float I_corr = r.I * r.image_scale_corr;
            const float sigma_corr = r.sigma * r.image_scale_corr;
            if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f)
                continue;

            const auto shell = shells.GetShell(r.d);
            if (!shell.has_value())
                continue;
            const int s = *shell;
            if (s >= 0 && s < n_shells) {
                acc[s].total_obs++;
                const auto key = generator(r).pack();
                const auto mit = merged_I.find(key);
                if (mit != merged_I.end()) {
                    auto &ra = rmeas_obs[key];
                    ra.sum_abs_dev += std::abs(static_cast<double>(I_corr) - mit->second);
                    ra.sum_I += I_corr;
                    ra.n++;
                    ra.shell = s;
                }
            }
        }
    }

    // R_meas per shell: sum over reflections of sqrt(n/(n-1)) * sum_i|I_i - <I>|, over sum of I_i.
    std::vector<double> rmeas_num(n_shells, 0.0), rmeas_den(n_shells, 0.0);
    double rmeas_num_all = 0.0, rmeas_den_all = 0.0;
    for (const auto &[key, ra]: rmeas_obs) {
        if (ra.n < 2 || ra.shell < 0 || ra.shell >= n_shells)
            continue;
        const double factor = std::sqrt(static_cast<double>(ra.n) / (ra.n - 1));
        rmeas_num[ra.shell] += factor * ra.sum_abs_dev;
        rmeas_den[ra.shell] += ra.sum_I;
        rmeas_num_all += factor * ra.sum_abs_dev;
        rmeas_den_all += ra.sum_I;
    }

    MergeStatistics out;
    out.shells.resize(n_shells);

    for (int s = 0; s < n_shells; ++s) {
        const auto &sa = acc[s];
        auto &ss = out.shells[s];

        ss.mean_one_over_d2 = shell_mean_1_d2[s];
        ss.d_min = shell_min_res[s];
        ss.d_max = s == 0 ? d_max_pad : shell_min_res[s - 1];
        ss.total_observations = sa.total_obs;
        ss.unique_reflections = sa.unique;
        ss.possible_unique_reflections = sa.possible;
        ss.mean_i_over_sigma = sa.n_i_over_sigma > 0
                                   ? sa.sum_i_over_sigma / sa.n_i_over_sigma
                                   : 0.0;

        ss.cc_half = sa.cc_half.GetCC();
        ss.cc_ref = sa.cc_ref.GetCC();
        ss.r_meas = rmeas_den[s] > 0.0 ? rmeas_num[s] / rmeas_den[s] : NAN;
    }

    auto &overall = out.overall;
    overall.d_min = d_min;
    overall.d_max = d_max;

    int all_possible = 0;
    int all_unique = 0;
    double sum_i_over_sigma = 0.0;
    int n_i_over_sigma = 0;


    for (const auto &sa: acc) {
        overall.total_observations += sa.total_obs;
        all_unique += sa.unique;
        all_possible += sa.possible;
        sum_i_over_sigma += sa.sum_i_over_sigma;
        n_i_over_sigma += sa.n_i_over_sigma;

    }

    overall.possible_unique_reflections = all_possible;
    overall.unique_reflections = all_unique;
    overall.mean_i_over_sigma = n_i_over_sigma > 0 ? sum_i_over_sigma / n_i_over_sigma : 0.0;
    overall.cc_half = cc_half_overall.GetCC();
    overall.cc_ref = cc_ref_overall.GetCC();
    overall.r_meas = rmeas_den_all > 0.0 ? rmeas_num_all / rmeas_den_all : NAN;

    return out;
}

std::ostream &operator<<(std::ostream &output, const MergeStatisticsShell &in) {
    double completeness = in.possible_unique_reflections > 0
        ? static_cast<double>(in.unique_reflections) / in.possible_unique_reflections * 100.0 : 0.0;
    double multiplicity = in.unique_reflections > 0
        ? static_cast<double>(in.total_observations) / in.unique_reflections : 0.0;

    output << fmt::format("{:8d} {:8d} {:8d} {:7.1f}% {:7.1f} {:8.1f} {:7.1f}% {:7.1f}% {:7.1f}%",
                in.total_observations,
                in.unique_reflections,
                in.possible_unique_reflections,
                completeness,
                multiplicity,
                in.mean_i_over_sigma,
                in.r_meas*100.0,
                in.cc_half*100.0,
                in.cc_ref*100.0);
    return output;
}

std::ostream &operator<<(std::ostream &output, const MergeStatistics &in) {
    output << std::endl;
    output << fmt::format("  {:>8s} {:>8s} {:>8s} {:>8s} {:>8s} {:>7s} {:>8s} {:>8s} {:>8s} {:>8s}",
        "d_min", "N_obs", "N_uniq", "N_possib", "Compl", "Mult", "<I/sig>", "R_meas", "CC1/2", "CCref")
    << std::endl;
    output << fmt::format("  {:->8s} {:->8s} {:->8s} {:->8s} {:->8s} {:->7s} {:->8s} {:->8s} {:->8s} {:->8s}",
        "", "", "", "", "", "", "", "", "", "") << std::endl;
    for (const auto &sh: in.shells) {
        if (sh.unique_reflections == 0)
            continue;
        output << fmt::format("  {:8.2f} ", sh.d_min);
        output << sh;
        output << std::endl;
    }
    output << fmt::format("  {:->8s} {:->8s} {:->8s} {:->8s} {:->8s} {:->7s} {:->8s} {:->8s} {:->8s} {:->8s}",
        "", "", "", "", "", "", "", "", "", "") << std::endl;

    output << fmt::format("  {:>8s} ", "Overall");
    output << in.overall;
    output << std::endl;
    if (std::isfinite(in.wilson_b) && in.wilson_b > 0.0)
        output << fmt::format("  Wilson B-factor estimate: {:.2f} A^2 (correlation {:.3f})",
                              in.wilson_b, in.wilson_b_correlation) << std::endl;
    output << std::endl;
    return output;
}