Jungfraujoch/image_analysis/scale_merge/Merge.h

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

#pragma once

#include <algorithm>
#include <cmath>
#include <map>
#include <random>
#include <unordered_map>
#include <vector>

#include "../../common/Logger.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/Reflection.h"
#include "../IntegrationOutcome.h"

#include "HKLKey.h"

// The error model splits a reflection's variance into a statistical part (a*sigma^2, which averages
// down with multiplicity) and a systematic part ((b*I)^2 - absorption, beam flicker, partiality,
// detector non-uniformity - correlated across a reflection's repeats). Inverse-variance merging
// (sigma = 1/sqrt(sum_w)) wrongly divides BOTH by the multiplicity, so high-multiplicity reflections
// get an unphysically small merged sigma (merged I/sigma far above ISa). Floor the merged sigma at
// b*|I| so the systematic term survives the merge and ISa = 1/b stays the asymptotic I/sigma ceiling.
// error_model_b <= 0 (no active error model) leaves the sigma unchanged.
inline float SigmaWithSystematicFloor(double inv_variance_sigma, float merged_I, double error_model_b) {
    const auto floor = static_cast<float>(error_model_b * std::abs(static_cast<double>(merged_I)));
    return std::max(static_cast<float>(inv_variance_sigma), floor);
}

struct MergeStatisticsShell {
    float d_min = 0.0f;
    float d_max = 0.0f;
    float mean_one_over_d2 = 0;

    int total_observations = 0;
    int unique_reflections = 0;
    int possible_unique_reflections = 0;

    double mean_i_over_sigma = 0.0;

    double cc_half = 0.0f;
    double cc_ref = NAN;

    // Redundancy-independent merging R-factor (Diederichs & Karplus 1997), computed over the
    // observations that enter the merge: R_meas = sum_hkl sqrt(n/(n-1)) sum_i|I_i-<I>| / sum I_i.
    double r_meas = NAN;
};

struct MergeStatistics {
    std::vector<MergeStatisticsShell> shells;
    MergeStatisticsShell overall;
};


std::ostream &operator<<(std::ostream &output, const MergeStatisticsShell &in);
std::ostream &operator<<(std::ostream &output, const MergeStatistics &in);

struct MergeAccum {
    int32_t h = 0;
    int32_t k = 0;
    int32_t l = 0;
    float d = NAN;
    double sum_wI = 0.0;
    double sum_w = 0.0;
    double sum_wI_half[2] = {0.0, 0.0};
    double sum_w_half[2] = {0.0, 0.0};
    size_t n_half[2] = {0, 0};
};

class MergeOnTheFly {
    mutable std::mutex merged_mutex;

    const int space_group_number = 1;

    ScalingSettings scaling_settings;
    IndexingSettings indexing_settings;

    std::optional<UnitCell> reference_cell;
    std::optional<double> high_resolution_limit;
    std::optional<double> image_cc_limit;
    double min_partiality = 0.02;

    // When set, ice-ring-flagged reflections are left out of this merge. Used for the P1 pass whose
    // merged intensities feed the space-group search and the error model - those model fits must not
    // see the ice-contaminated intensities. The final in-symmetry merge keeps them (for completeness).
    bool exclude_ice_rings = false;

    // Ice rings (indices into ICE_RING_RES_A) to exclude from this merge because their merged CC1/2
    // collapsed relative to their resolution neighbours - ice decorrelated them. Empty = none masked.
    std::vector<char> masked_ice_rings;
    float mask_ice_half_width_q = 0.03f;

    HKLKeyGenerator generator;

    std::map<uint64_t, MergeAccum> accumulator;

    // Global error model (XDS form): sigma_corr^2 = a*sigma^2 + (b*<I>)^2. a rescales the
    // (under-estimated) counting variance; the (b*<I>)^2 term adds the intensity-
    // proportional systematic error that counting statistics miss, so strong reflections
    // are no longer over-weighted. ISa = 1/b is the asymptotic I/sigma. Refined from the
    // scatter of symmetry equivalents (RefineErrorModel); identity until then.
    bool   error_model_active = false;
    double error_model_a = 1.0;
    double error_model_b = 0.0;
    double error_model_chi2 = 0.0;   // achieved median reduced chi^2 (~1.0 = honestly calibrated sigmas)
    // The (b*I)^2 term uses the reflection's *mean* intensity (constant over its
    // observations), so it inflates sigma without biasing the inverse-variance weights -
    // using the per-observation I_i instead would over-weight down-fluctuated points.
    std::unordered_map<uint64_t, float> error_model_mean_I;
    [[nodiscard]] float CorrectedSigma(float I_corr, float sigma_corr, uint64_t hkl_key) const;

    // Optional per-observation outlier rejection: drop observations whose corrected
    // intensity lies more than reject_nsigma error-model sigmas from the reflection's
    // *median* (a robust centre). The error-model sigma already captures the genuine
    // (e.g. partiality) scatter, so this removes only the tail beyond it - zingers,
    // overlaps, mis-indexed frames - not good partials. Populated by RefineErrorModel.
    bool   reject_outliers = false;
    double reject_nsigma = 6.0;
    std::unordered_map<uint64_t, float> reject_median_I;
    size_t reject_count = 0;

    bool Mask(const IntegrationOutcome &outcome, bool cc_mask);
    [[nodiscard]] bool IsMaskedRing(const Reflection &r) const;
public:
    MergeOnTheFly(const DiffractionExperiment &x);
    MergeOnTheFly& ReferenceCell(const std::optional<UnitCell> &cell);
    MergeOnTheFly& ExcludeIceRings(bool input) { exclude_ice_rings = input; return *this; }
    MergeOnTheFly& MaskIceRings(std::vector<char> masked, float half_width_q) {
        masked_ice_rings = std::move(masked); mask_ice_half_width_q = half_width_q; return *this;
    }

    // Fit the global error model from the spread of symmetry-equivalent observations.
    // Call once before merging; AddImage then applies it.
    void RefineErrorModel(const std::vector<IntegrationOutcome> &outcomes);
    [[nodiscard]] bool   ErrorModelActive() const { return error_model_active; }
    [[nodiscard]] double ErrorModelA() const { return error_model_a; }
    [[nodiscard]] double ErrorModelB() const { return error_model_b; }
    [[nodiscard]] double ErrorModelChi2() const { return error_model_chi2; }

    // Outlier rejection (driven by ScalingSettings::GetOutlierRejectNsigma) reports its count.
    [[nodiscard]] size_t RejectedCount() const { return reject_count; }

    // image_id is the image's stable identity (its index in the outcomes vector). The CC1/2 half-set
    // is a deterministic hash of it, so the split is reproducible run-to-run and independent of the
    // order (or threading) of AddImage calls - not a draw from a shared RNG in call order.
    void AddImage(const IntegrationOutcome& outcome, int64_t image_id, bool cc_mask = false);

    // Per-crystal CC1/2-delta rejection (CrystFEL deltaCChalf): returns a per-image flag
    // marking images whose removal would raise CC1/2 by a low-side outlier amount
    // (deltaCChalf < mean - nsigma*stddev). Skip the flagged images when merging.
    [[nodiscard]] std::vector<char> DeltaCChalfReject(const std::vector<IntegrationOutcome> &outcomes,
                                                      double nsigma) const;

    // d_min_override, when set, is the effective high-resolution limit for the shell table (used for
    // the automatic resolution cutoff computed by the caller); otherwise the manual
    // ScalingSettings high-resolution limit stands. The number of shells is ScalingSettings::ReportShellCount.
    MergeStatistics MergeStats(const std::vector<MergedReflection> &merged,
                               const std::vector<IntegrationOutcome> &reflections,
                               const std::vector<MergedReflection> &reference = {},
                               std::optional<double> d_min_override = std::nullopt);

    std::vector<MergedReflection> ExportReflections();
};

std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
                                       const std::vector<IntegrationOutcome> &reflections,
                                       bool mask = false);