Jungfraujoch/image_analysis/scale_merge/RotationScaleMerge.h

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

#pragma once

#include <cstdint>
#include <optional>
#include <vector>

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

#include "Merge.h"        // MergedReflection, MergeStatistics
#ifdef JFJOCH_USE_CUDA
#include <memory>
#include "RotationScaleMergeGPU.h"
#endif

// Dedicated, allocate-once scale+combine+merge for rotation data (the -P rot3d path).
//
// This is a distinct, faster path from ScaleOnTheFly + MergeAll + CombineRotationObservations +
// MergeOnTheFly. Those rebuild a std::map keyed by hkl on *every* scaling iteration and every merge
// (7-14 map rebuilds per space-group pass), which dominates the offline wall clock. Here the per-frame
// partial observations are ingested ONCE into flat vectors; the hkl->ASU grouping is computed once per
// space group (by a sort, not a map) and reused across all scaling iterations; every hot step is a flat
// loop over those vectors, so it also maps directly onto CUDA kernels (segmented reduction + per-frame
// solve). CC1/2 and the per-image CC are computed once at the end, not every iteration.
//
// It reproduces the numerics of the CPU pipeline exactly (same robust IRLS per-frame G, same 3D combine,
// same XDS-order scale-fulls, same global error model, same merge statistics) - the speed-up is purely
// from the data layout, not from cutting corners. It is used only for the self-scaling rotation case
// with per-image G (Rotation partiality, no B refinement, no external reference, no absorption surface);
// stills, B-factor refinement, reference scaling and the absorption surface stay on the classic path.
class RotationScaleMerge {
public:
    struct Result {
        std::vector<MergedReflection> merged;
        MergeStatistics statistics;
        double isa = 0.0;   // 1/b of the fitted error model (0 if the model stayed at identity)
    };

    // experiment: read live (its space group is changed by the caller between Run() calls).
    // partial_outcomes: the per-frame partials; the final per-frame scale (G, CC, mosaicity) is written
    //   back onto them so the offline per-image scaling table is still exported.
    // reference_cell: the consensus cell (for the completeness count and the cell-consistency mask).
    RotationScaleMerge(const DiffractionExperiment &experiment,
                       std::vector<IntegrationOutcome> &partial_outcomes,
                       std::optional<UnitCell> reference_cell,
                       int scaling_iterations,
                       float ice_ring_half_width_q,
                       size_t nthreads,
                       Logger &logger,
                       std::string observation_dump_path = {});

    // Copy the per-frame partials into the flat buffers. Call once before the first Run().
    void Ingest();

    // Scale (per-frame G) -> smooth G -> 3D combine -> scale fulls -> merge -> error model -> statistics
    // for the space group currently set on the experiment, reusing the ingested buffers.
    // for_search: the de-novo P1 pass whose merged intensities feed the space-group search - ice-ring
    //   reflections are dropped from the merge and the error model (kept otherwise, for completeness).
    // masked_ice_rings: rings (indices into ICE_RING_RES_A) to drop from the final merge; empty = none.
    Result Run(bool for_search, const std::vector<char> &masked_ice_rings = {});

private:
    // One integrated observation - a per-frame partial during scaling/combine, or a combined full during
    // scale-fulls/merge. Flat (not nested per image); a POD so the arrays translate straight to CUDA.
    struct Obs {
        int32_t h, k, l;
        float I, sigma, d, rlp, partiality, zeta, delta_phi, bkg;
        float image_number;   // fractional frame position (for 3D-combine contiguity)
        int32_t frame;        // index of the outcome whose per-frame scale G applies to this obs
        uint8_t on_ice;
        float corr;           // image_scale_corr (working; updated by scaling)
        int32_t group;        // dense ASU-group id for the current space group; <0 = never mergeable
    };

    const DiffractionExperiment &x;
    std::vector<IntegrationOutcome> &partials_out;   // written back at the end of scaling
    std::optional<UnitCell> reference_cell;
    size_t nthreads;
    Logger &logger;
    std::string observation_dump_path;

    // Fixed settings snapshot (read once in the ctor).
    int n_frames = 0;
    double min_partiality = 0.02;
    std::optional<double> d_min_limit;
    bool merge_friedel = true;
    double capture_uncertainty_coeff = 0.0;
    double reject_nsigma = 0.0;
    bool reject_outliers = false;
    double rfree_fraction = 0.0;
    int scaling_iter = 3;
    bool scale_fulls = true;
    double mosaicity_deg = 0.1;
    float ice_half_width_q = 0.0f;

    // Flat buffers, allocated once by Ingest() and reused across Run() calls.
    std::vector<Obs> partials;                    // all per-frame partials, grouped by frame
    std::vector<int32_t> frame_start, frame_count; // CSR ranges of `partials` per frame
    std::vector<uint8_t> frame_cell_ok;           // per-frame cell-consistency mask (1 = kept)
    std::vector<double> g_partial;                // per-frame partial scale G

    // Raw-hkl ordering, built ONCE by Ingest and reused: `perm` lists partial indices sorted by
    // (raw h,k,l, image_number); each distinct raw hkl is a contiguous run [rawrun_start, +count) of it.
    // The expensive sort happens once here, so per-pass combine (event split) and ASU grouping are linear.
    std::vector<int32_t> perm;
    std::vector<int32_t> rawrun_start, rawrun_count;
    std::vector<int32_t> rawrun_h, rawrun_k, rawrun_l;
    std::vector<float>   rawrun_d;                 // representative resolution per raw hkl
    std::vector<int32_t> rawrun_group;            // dense ASU-group id per raw hkl (<0 = absent/out of range)

    std::vector<Obs> fulls;                       // combined fulls (rebuilt each Run), sorted by frame
    std::vector<int32_t> fulls_frame_start, fulls_frame_count;  // CSR ranges of `fulls` per frame
    std::vector<double> g_full;                   // per-frame scale on the fulls

    // Set by FitPerFrameG: which frames were fitted this call (so corr/G is updated only there).
    std::vector<uint8_t> frame_scaled_scratch;

    // Per-frame mosaicity smoothed in frame order (deterministic); used to recompute partiality and
    // written back for the per-image scaling table. Empty if there is no per-frame mosaicity.
    std::vector<float> mos_smooth;

    // Working per-group arrays (sized to the current group count; reused).
    std::vector<int32_t> group_h, group_k, group_l;

#ifdef JFJOCH_USE_CUDA
    // GPU engine for the partial-scaling loop (segmented reduce + per-frame IRLS + corr update). Null /
    // inactive when no GPU; the CPU loops are the fallback. Built in Ingest.
    std::unique_ptr<RotationScaleMergeGPU> gpu_;
    bool gpu_active_ = false;
    // Phase-2 GPU combine (partials->fulls on the device). Opt-in for now (JFJOCH_RSM_GPU_COMBINE):
    // it is validated bit-for-bit against the CPU combine and deterministic, but until the fulls stay
    // resident for the scale-fulls + merge stages the fulls round-trip makes it no faster than the CPU
    // path, so it is not on by default. See RotationScaleMergeGPU::Combine.
    bool gpu_combine_ = false;
#endif

    // --- helpers (each a flat pass; see the .cpp) ---
    // Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl
    // runs by their ASU key - one gemmi ASU reduction per distinct raw hkl, not per observation. Fills
    // rawrun_group, the group_h/k/l representative tables, and partials[].group; returns the group count.
    int ComputeAsuGroups(const HKLKeyGenerator &key_generator);

    // Inverse-variance per-group mean of I*corr over `obs` (the merge reference). exclude_ice/masked drop
    // those reflections (used for the error-model/merge means, not the scaling reference).
    void ReduceGroupMeans(const std::vector<Obs> &obs, int n_groups,
                          bool exclude_ice, const std::vector<char> &masked_ice_rings,
                          std::vector<double> &out_mean) const;

    // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the
    // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity.
    void FitPerFrameG(std::vector<Obs> &obs, const std::vector<int32_t> &fstart,
                      const std::vector<int32_t> &fcount, const std::vector<double> &group_mean_in,
                      bool unity, std::vector<double> &g);

    // corr = rlp / (partiality * G[frame]); leaves corr unchanged for frames that could not be fit.
    void UpdateCorr(std::vector<Obs> &obs, const std::vector<double> &g,
                    const std::vector<uint8_t> &frame_scaled) const;

    void SmoothG(std::vector<Obs> &obs, std::vector<double> &g, int window) const;

    // Smooth per-frame mosaicity in frame order and recompute each partial's partiality from it, so the
    // per-frame partials of one rocking event tile the curve consistently (they sum toward 1) before the
    // 3D combine. Deterministic (frame order); replaces the old arrival-order mosaicity moving average
    // that prediction applied. SG-independent, so done once in Ingest.
    void SmoothMosaicityAndPartiality();

    void Combine();   // partials -> fulls (CPU)

    // Sort `fulls` by peak frame and (re)build fulls_frame_start/count (the per-frame CSR the scale-fulls
    // step slices). Shared by the CPU Combine tail and the GPU combine path.
    void SortFullsByFrame();

    // Per-frame CC vs the partial merge reference, then write G/CC/mosaicity back onto the partials
    // (once, at the end of partial scaling) so the offline per-image scaling table is still exported.
    void FinalizePerFrameScale(int n_groups, const std::vector<double> &partial_group_mean,
                               const std::vector<uint8_t> &frame_scaled);

    // Error model + merge + statistics over the fulls (the last stage). n_groups is the fulls group count.
    Result MergeAndStats(int n_groups, bool for_search, const std::vector<char> &masked_ice_rings);
};