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

#include "RotationScaleMergeGPU.h"

#include <algorithm>
#include <cmath>
#include <string>
#include <cuda_runtime.h>

#include "../indexing/CUDAMemHelpers.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"

namespace {
    constexpr int BLK = 256;
    constexpr int MIN_REFLECTIONS = 20;

    __device__ __forceinline__ double SafeInvD(double x, double fallback) {
        return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback;
    }

    // Block reduction of a double, deterministic for a fixed thread->element mapping (fixed order,
    // no atomics). Returns the sum on thread 0; `s` is BLK doubles of shared scratch.
    __device__ double BlockReduceSum(double v, double *s) {
        const int t = threadIdx.x;
        s[t] = v;
        __syncthreads();
        for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
            if (t < stride) s[t] += s[t + stride];
            __syncthreads();
        }
        return s[0];
    }

    // One thread per ASU group (grid-stride): inverse-variance mean of I*corr over the group's contiguous,
    // fixed-order segment of group_perm. Groups are small (avg tens of obs), so a whole block per group
    // wastes threads; summing each group in one thread avoids launch/sync overhead and stays deterministic
    // (fixed group_perm order). Matches the CPU ReduceGroupMeans filter (no ice/cell mask - scaling ref).
    __global__ void ReduceGroupMeansKernel(int n_groups, double min_partiality,
                                           const int32_t *__restrict__ group_perm,
                                           const int32_t *__restrict__ group_start,
                                           const int32_t *__restrict__ group_count,
                                           const float *__restrict__ I, const float *__restrict__ sigma,
                                           const float *__restrict__ partiality,
                                           const float *__restrict__ corr,
                                           double *__restrict__ group_mean) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) {
            const int lo = group_start[g], hi = group_start[g] + group_count[g];
            double sw = 0.0, swI = 0.0;
            for (int p = lo; p < hi; ++p) {
                const int i = group_perm[p];
                const float c = corr[i];
                if (!(c > 0.0f) || !isfinite(c)) continue;
                if (partiality[i] < min_partiality) continue;
                const float I_corr = I[i] * c;
                const float sigma_corr = sigma[i] * c;
                if (!isfinite(I_corr) || !isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
                const double w = 1.0 / (double(sigma_corr) * sigma_corr);
                sw += w;
                swI += w * I_corr;
            }
            group_mean[g] = sw > 0.0 ? swI / sw : NAN;
        }
    }

    // Per-observation scale-fit coefficient (rotation model) and accept flag, recomputed each scaling
    // iteration once the group means are known. coeff = partiality * (1/rlp) * mean[group].
    __global__ void PrepScaleObsKernel(int n_obs, const int32_t *__restrict__ group,
                                       const float *__restrict__ partiality, const float *__restrict__ rlp,
                                       const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
                                       const double *__restrict__ group_mean,
                                       float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n_obs) return;
        const int g = group[i];
        bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f;
        double mean = 0.0;
        if (ok) {
            mean = group_mean[g];
            ok = isfinite(mean);
        }
        sco_ok[i] = ok ? 1 : 0;
        sco_coeff[i] = ok ? float(double(partiality[i]) * SafeInvD(rlp[i], 1.0) * mean) : 0.0f;
    }

    // One block per frame: robust per-frame scale G by IRLS (Cauchy), over the frame's contiguous obs.
    // Identical objective to the CPU SolveScaleIRLS. Leaves g/scaled untouched for under-populated frames.
    // `perm` (null for the partials, whose arrays are already frame-contiguous) maps a position in the
    // frame's [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a
    // frame-grouping permutation without physically reordering the fulls arrays.
    __global__ void FitPerFrameGKernel(int n_frames, double robust_k,
                                       const int32_t *__restrict__ frame_start,
                                       const int32_t *__restrict__ frame_count,
                                       const float *__restrict__ I, const float *__restrict__ sigma,
                                       const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok,
                                       const int32_t *__restrict__ perm,
                                       double *__restrict__ g, uint8_t *__restrict__ scaled) {
        const int f = blockIdx.x;
        if (f >= n_frames) return;
        const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
        __shared__ double sh[BLK];

        // count accepted
        long cnt_local = 0;
        for (int i = lo + threadIdx.x; i < hi; i += blockDim.x)
            if (sco_ok[perm ? perm[i] : i]) ++cnt_local;
        const double cnt = BlockReduceSum(double(cnt_local), sh);
        __shared__ double s_cnt;
        if (threadIdx.x == 0) s_cnt = cnt;
        __syncthreads();
        if (s_cnt < MIN_REFLECTIONS) return;   // leave g[f]/scaled[f] as-is

        const double k2 = robust_k * robust_k;

        // seed: plain weighted-LS ratio (robust weight = 1)
        double num = 0.0, den = 0.0;
        for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
            const int a = perm ? perm[i] : i;
            if (!sco_ok[a]) continue;
            const double coeff = sco_coeff[a];
            const double w = SafeInvD(sigma[a], 1.0);
            const double w2 = w * w;
            num += w2 * coeff * double(I[a]);
            den += w2 * coeff * coeff;
        }
        double tnum = BlockReduceSum(num, sh); __syncthreads();
        double tden = BlockReduceSum(den, sh);
        __shared__ double s_G;
        if (threadIdx.x == 0) {
            double G = tden > 0.0 ? tnum / tden : NAN;
            s_G = isfinite(G) ? fmax(0.0, G) : 1.0;
        }
        __syncthreads();
        if (!(s_G >= 0.0)) { if (threadIdx.x == 0) { g[f] = 1.0; scaled[f] = 1; } return; }

        for (int iter = 0; iter < 30; ++iter) {
            const double G = s_G;
            num = 0.0; den = 0.0;
            for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
                const int a = perm ? perm[i] : i;
                if (!sco_ok[a]) continue;
                const double coeff = sco_coeff[a];
                const double w = SafeInvD(sigma[a], 1.0);
                const double w2 = w * w;
                const double res = w * (G * coeff - double(I[a]));
                const double rw = 1.0 / (1.0 + res * res / k2);
                num += rw * w2 * coeff * double(I[a]);
                den += rw * w2 * coeff * coeff;
            }
            tnum = BlockReduceSum(num, sh); __syncthreads();
            tden = BlockReduceSum(den, sh);
            if (threadIdx.x == 0) {
                const double G_next = tden > 0.0 ? tnum / tden : NAN;
                if (isfinite(G_next)) {
                    const double Gn = fmax(0.0, G_next);
                    s_cnt = (fabs(Gn - G) <= 1e-7 * fmax(Gn, 1.0)) ? 1.0 : 0.0; // reuse s_cnt as "converged"
                    s_G = Gn;
                } else {
                    s_cnt = 1.0;
                }
            }
            __syncthreads();
            if (s_cnt != 0.0) break;
        }
        if (threadIdx.x == 0) { g[f] = s_G; scaled[f] = 1; }
    }

    // One block per frame: Pearson CC of (I*corr) vs the merged group mean over the frame's partials,
    // == FinalizePerFrameScale's per-frame loop. Diagnostic only (per-image scaling table), so the tree
    // reduction's ~ulp difference from the CPU is immaterial; deterministic run-to-run.
    __global__ void PerFrameCCKernel(int n_frames, double min_partiality,
                                     const int32_t *__restrict__ frame_start,
                                     const int32_t *__restrict__ frame_count,
                                     const float *__restrict__ I, const float *__restrict__ sigma,
                                     const float *__restrict__ partiality, const float *__restrict__ corr,
                                     const uint8_t *__restrict__ on_ice, const int32_t *__restrict__ group,
                                     const double *__restrict__ group_mean,
                                     double *__restrict__ cc_out, int64_t *__restrict__ cc_n_out) {
        const int f = blockIdx.x;
        if (f >= n_frames) return;
        const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
        __shared__ double sh[BLK];
        double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
        long nl = 0;
        for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
            if (on_ice[i]) continue;
            const int g = group[i];
            if (g < 0) continue;
            if (partiality[i] < min_partiality) continue;
            const float c = corr[i];
            if (!isfinite(I[i]) || !isfinite(c) || !(c > 0.0f)) continue;
            if (!isfinite(sigma[i]) || !(sigma[i] > 0.0f)) continue;
            const double mean = group_mean[g];
            if (!isfinite(mean)) continue;
            const double img = double(I[i]) * c;
            sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; ++nl;
        }
        const double tsx = BlockReduceSum(sx, sh); __syncthreads();
        const double tsy = BlockReduceSum(sy, sh); __syncthreads();
        const double tsx2 = BlockReduceSum(sx2, sh); __syncthreads();
        const double tsy2 = BlockReduceSum(sy2, sh); __syncthreads();
        const double tsxy = BlockReduceSum(sxy, sh); __syncthreads();
        const double tn = BlockReduceSum(double(nl), sh);
        if (threadIdx.x == 0) {
            cc_out[f] = NAN; cc_n_out[f] = 0;
            if (tn >= MIN_REFLECTIONS) {
                const double cov = tsxy - tsx * tsy / tn;
                const double vx = tsx2 - tsx * tsx / tn;
                const double vy = tsy2 - tsy * tsy / tn;
                if (vx > 0.0 && vy > 0.0) { cc_out[f] = cov / sqrt(vx * vy); cc_n_out[f] = int64_t(tn); }
            }
        }
    }

    // SmoothG corr adjust: corr[i] *= ratio[frame[i]] for frames flagged apply, in double then stored
    // as float - matching CPU SmoothG's `corr = float(corr * (g/g_smooth))`. Grid-stride, resident corr.
    __global__ void SmoothCorrKernel(int n_obs, const int32_t *__restrict__ frame,
                                     const uint8_t *__restrict__ apply, const double *__restrict__ ratio,
                                     float *__restrict__ corr) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            if (!apply[frame[i]]) continue;
            const float c = corr[i];
            if (isfinite(c)) corr[i] = float(double(c) * ratio[frame[i]]);
        }
    }

    // corr = rlp / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
    __global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame,
                                     const float *__restrict__ rlp, const float *__restrict__ partiality,
                                     const double *__restrict__ g, const uint8_t *__restrict__ scaled,
                                     float *__restrict__ corr) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            const int f = frame[i];
            if (!scaled[f]) continue;
            const double denom = double(partiality[i]) * g[f];
            corr[i] = (isfinite(double(rlp[i])) && isfinite(denom) && denom > 0.0)
                          ? float(rlp[i] / denom) : NAN;
        }
    }

    // std::max / std::min return (a<b)?b:a and (b<a)?b:a - reproduce that exactly (NOT fmax/fmin, which
    // differ on NaN) so the combine matches the CPU path bit-for-bit on the same inputs.
    __device__ __forceinline__ double Dmax(double a, double b) { return (a < b) ? b : a; }
    __device__ __forceinline__ double Dmin(double a, double b) { return (b < a) ? b : a; }
    __device__ __forceinline__ float  Fmax(float a, float b)   { return (a < b) ? b : a; }

    constexpr float COMBINE_MAX_FRAME_GAP = 2.0f;   // == RotationScaleMerge::MAX_FRAME_GAP

    // A partial is usable for the combine iff its corr and (I, sigma) are finite with corr>0, sigma>0.
    __device__ __forceinline__ bool CombineUsable(int i, const float *I, const float *sigma,
                                                  const float *corr) {
        const float c = corr[i];
        if (!(c > 0.0f) || !isfinite(c)) return false;
        return isfinite(I[i]) && isfinite(sigma[i]) && sigma[i] > 0.0f;
    }

    // All device pointers + scalars the combine kernels need, passed by value.
    struct CombineParams {
        int n_runs;
        double min_partiality, capture_uncertainty_coeff, min_captured_fraction;
        const float *I, *sigma, *corr, *partiality, *bkg, *image_number, *d, *px, *py;
        const int32_t *frame;
        const uint8_t *on_ice;
        const int32_t *perm, *rr_start, *rr_count, *rr_h, *rr_k, *rr_l, *rr_group;
        int32_t *rr_nevents, *rr_nusable;   // count pass outputs
        const int32_t *rr_offset;           // emit pass: per-run base offset into the fulls arrays
        int32_t *f_h, *f_k, *f_l, *f_frame, *f_group;
        float *f_I, *f_sigma, *f_d, *f_img, *f_px, *f_py;
        uint8_t *f_on_ice;
    };

    // One thread's raw-hkl run: split its usable partials (already in image_number order within the run)
    // into rocking events, and for each event pool background, seed F and run the 3-iter de-biased Poisson
    // reweight - the exact objective of RotationScaleMerge::Combine::process_rawrun. When Emit, write the
    // resulting full at rr_offset[r] + (event index); otherwise just count the emitted events. Both modes
    // run the identical accept test, so the count pass predicts the emit pass exactly.
    template <bool Emit>
    __device__ void CombineRawRun(int r, const CombineParams &p) {
        const int lo = p.rr_start[r];
        const int hi = lo + p.rr_count[r];
        const int group = p.rr_group[r];

        int n_emit = 0;
        int cursor = lo;
        while (cursor < hi) {
            while (cursor < hi && !CombineUsable(p.perm[cursor], p.I, p.sigma, p.corr)) ++cursor;
            if (cursor >= hi) break;
            const int ev_start = cursor;             // first usable position of the event
            int ev_end = cursor;                     // last usable position (inclusive), extended below
            float last_img = p.image_number[p.perm[cursor]];
            int probe = cursor + 1;
            while (probe < hi) {
                while (probe < hi && !CombineUsable(p.perm[probe], p.I, p.sigma, p.corr)) ++probe;
                if (probe >= hi) break;
                const float img = p.image_number[p.perm[probe]];
                if (img - last_img > COMBINE_MAX_FRAME_GAP) break;
                last_img = img;
                ev_end = probe;
                ++probe;
            }
            cursor = ev_end + 1;

            // Pass A: pooled background = mean of the event members' finite backgrounds.
            double pooled_bkg = 0.0;
            int n_pool = 0;
            for (int m = ev_start; m <= ev_end; ++m) {
                const int i = p.perm[m];
                if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
                const float b = p.bkg[i];
                if (isfinite(b)) { pooled_bkg += b; ++n_pool; }
            }
            pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;

            auto pooled_I = [&](int i) -> double {
                const double n_bkg = Dmax(0.0, double(p.sigma[i]) * p.sigma[i] - p.I[i])
                                     / Fmax(p.bkg[i], 1.0f);
                return double(p.I[i]) + n_bkg * (double(p.bkg[i]) - pooled_bkg);
            };

            // Pass B: seed F (inverse-variance mean of pooled_I*corr), plus peak / d / on_ice / partiality.
            double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
            float d = NAN;
            const int first = p.perm[ev_start];
            int peak_outcome = p.frame[first];
            float peak_frame = p.image_number[first];
            float peak_px = p.px[first], peak_py = p.py[first];
            float peak_partiality = -1.0f;
            const bool on_ice = p.on_ice[first];
            for (int m = ev_start; m <= ev_end; ++m) {
                const int i = p.perm[m];
                if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
                const double sigma_corr = double(p.sigma[i]) * p.corr[i];
                const double w = 1.0 / (sigma_corr * sigma_corr);
                sum_w += w;
                sum_wI += w * pooled_I(i) * p.corr[i];
                sum_partiality += p.partiality[i];
                if (p.partiality[i] > peak_partiality) {
                    peak_partiality = p.partiality[i];
                    peak_outcome = p.frame[i];
                    peak_frame = p.image_number[i];
                    peak_px = p.px[i]; peak_py = p.py[i];
                }
                if (!isfinite(d) && isfinite(p.d[i]) && p.d[i] > 0.0f) d = p.d[i];
            }
            double F = sum_wI / sum_w;

            // Pass C: 3 de-biased Poisson reweights (variance = bkg part + corr*max(0,F)).
            for (int iter = 0; iter < 3; ++iter) {
                sum_w = 0.0; sum_wI = 0.0;
                for (int m = ev_start; m <= ev_end; ++m) {
                    const int i = p.perm[m];
                    if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
                    const double corr = p.corr[i];
                    const double I_corr = pooled_I(i) * corr;
                    const double sigma_corr = double(p.sigma[i]) * corr;
                    const double bkg_var = sigma_corr * sigma_corr - corr * I_corr;
                    double var = Dmax(0.0, bkg_var) + corr * Dmax(0.0, F);
                    if (!(var > 0.0)) var = sigma_corr * sigma_corr;
                    const double w = 1.0 / var;
                    sum_w += w;
                    sum_wI += w * I_corr;
                }
                F = sum_wI / sum_w;
            }

            if (sum_w <= 0.0 || sum_partiality < p.min_partiality
                             || sum_partiality < p.min_captured_fraction)
                continue;

            double sigma_full = 1.0 / sqrt(sum_w);
            if (p.capture_uncertainty_coeff > 0.0) {
                const double frac = Dmin(1.0, sum_partiality);
                const double extra = p.capture_uncertainty_coeff * (1.0 - frac) * Dmax(0.0, F);
                sigma_full = sqrt(sigma_full * sigma_full + extra * extra);
            }

            if (Emit) {
                const int o = p.rr_offset[r] + n_emit;
                p.f_h[o] = p.rr_h[r]; p.f_k[o] = p.rr_k[r]; p.f_l[o] = p.rr_l[r];
                p.f_I[o] = float(F);
                p.f_sigma[o] = float(sigma_full);
                p.f_d[o] = d;
                p.f_img[o] = peak_frame;
                p.f_px[o] = peak_px; p.f_py[o] = peak_py;
                p.f_frame[o] = peak_outcome;
                p.f_on_ice[o] = on_ice ? 1 : 0;
                p.f_group[o] = group;
            }
            ++n_emit;
        }

        if (!Emit) {
            p.rr_nevents[r] = n_emit;
            int n_usable = 0;
            for (int m = lo; m < hi; ++m)
                if (CombineUsable(p.perm[m], p.I, p.sigma, p.corr)) ++n_usable;
            p.rr_nusable[r] = n_usable;
        }
    }

    template <bool Emit>
    __global__ void CombineKernel(CombineParams p) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < p.n_runs; r += gridDim.x * blockDim.x)
            CombineRawRun<Emit>(r, p);
    }

    template <typename T>
    __global__ void FillKernel(T *p, int n, T v) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v;
    }

    // ===== error-model + merge reductions over the resident, scaled fulls (mirror MergeAndStats) =====

    // Device copy of common/Definitions.h IceRingIndex (only consulted when the merge has masked rings).
    __device__ __forceinline__ int IceRingIndexDev(float d, float hw) {
        if (!(d > 0.0f)) return -1;
        const float two_pi = 6.283185307f;
        const float q = two_pi / d;
        const float rings[11] = {3.895f, 3.661f, 3.438f, 2.667f, 2.249f, 2.068f,
                                 1.947f, 1.916f, 1.882f, 1.719f, 1.522f};
        for (int i = 0; i < 11; ++i)
            if (fabsf(q - two_pi / rings[i]) < hw) return i;
        return -1;
    }

    // Deterministic CC1/2 half from the frame index (splitmix64), matching Merge.cpp / RSM HalfForImage.
    __device__ __forceinline__ int HalfForImageDev(long long image_id) {
        unsigned long long z = (unsigned long long)image_id + 0x9e3779b97f4a7c15ULL;
        z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
        z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
        z = z ^ (z >> 31);
        return (int)(z & 1ULL);
    }

    struct MergeParams {
        int n_groups;
        double min_partiality, error_model_a, error_model_b, reject_nsigma;
        float ice_hw;
        int for_search, n_masked, error_model_active, reject_outliers;
        const float *I, *sigma, *corr, *partiality, *d, *reject_median;
        const int32_t *group, *frame;
        const uint8_t *on_ice, *frame_cell_ok, *masked;
        const int32_t *gperm, *gstart, *gcount;
        const double *em_mean, *merged_I;
        // outputs
        double *sw, *swI, *em_mean_out;
        int32_t *cnt;
        double *s2, *I2, *dev2;
        uint8_t *valid;
        double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_d;
        int32_t *a_nh0, *a_nh1, *a_rejected;
        double *r_absdev, *r_sumI;
        int32_t *r_n, *r_nusable;
        uint8_t *rejected_obs;   // per-full: set by MergeAccum (outlier-rejected), read by MergeRmeas
    };

    // A full passes the merge / error-model filter (mirrors MergeAndStats::usable_merge).
    __device__ __forceinline__ bool MergeUsable(int i, const MergeParams &p) {
        const int g = p.group[i];
        if (g < 0) return false;
        if (!p.frame_cell_ok[p.frame[i]]) return false;
        const float c = p.corr[i];
        if (!(c > 0.0f) || !isfinite(c)) return false;
        if (p.for_search && p.on_ice[i]) return false;
        if (p.n_masked > 0) {
            const int ring = IceRingIndexDev(p.d[i], p.ice_hw);
            if (ring >= 0 && ring < p.n_masked && p.masked[ring]) return false;
        }
        if (p.partiality[i] < p.min_partiality) return false;
        const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
        return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
    }

    // The looser R_meas filter (no ice / masked-ring / for_search - Mask = cell only).
    __device__ __forceinline__ bool RmeasUsable(int i, const MergeParams &p) {
        const int g = p.group[i];
        if (g < 0) return false;
        if (!p.frame_cell_ok[p.frame[i]]) return false;
        const float c = p.corr[i];
        if (!(c > 0.0f) || !isfinite(c)) return false;
        if (p.partiality[i] < p.min_partiality) return false;
        const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
        return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
    }

    // One thread per group: inverse-variance sums over the group's usable fulls + the group mean (>=2).
    __global__ void MergeEmStatsKernel(MergeParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            double sw = 0.0, swI = 0.0;
            int cnt = 0;
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!MergeUsable(i, p)) continue;
                const double sigma_corr = double(p.sigma[i]) * p.corr[i];
                const double w = 1.0 / (sigma_corr * sigma_corr);
                sw += w; swI += w * (double(p.I[i]) * p.corr[i]); ++cnt;
            }
            p.sw[g] = sw; p.swI[g] = swI; p.cnt[g] = cnt;
            p.em_mean_out[g] = (cnt >= 2 && sw > 0.0) ? swI / sw : NAN;
        }
    }

    // One thread per full: the leverage-corrected error-model sample, or valid=0 if dropped.
    __global__ void MergeSamplesKernel(int n_obs, MergeParams p) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            p.valid[i] = 0;
            if (!MergeUsable(i, p)) continue;
            const int g = p.group[i];
            if (p.cnt[g] < 2) continue;
            const double mean = p.em_mean[g];
            if (!isfinite(mean)) continue;
            const double sigma_corr = double(p.sigma[i]) * p.corr[i];
            const double s2 = sigma_corr * sigma_corr;
            const double factor = 1.0 - (1.0 / s2) / p.sw[g];
            if (factor < 0.05) continue;
            const double resid = double(p.I[i]) * p.corr[i] - mean;
            p.s2[i] = s2; p.I2[i] = mean * mean; p.dev2[i] = resid * resid / factor; p.valid[i] = 1;
        }
    }

    // One thread per group: the merge accumulators (inv-var sums + deterministic half-sets), with the
    // error-model-corrected sigma. Mirrors MergeAndStats' merge loop (reject path stays on the host).
    __global__ void MergeAccumKernel(MergeParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            double swI = 0, sw = 0, swIh0 = 0, swIh1 = 0, swh0 = 0, swh1 = 0, dd = NAN;
            int nh0 = 0, nh1 = 0, rejected = 0;
            const float rmed = p.reject_median[g];
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!MergeUsable(i, p)) continue;
                const float I_corr = p.I[i] * p.corr[i];
                float sigma_corr = p.sigma[i] * p.corr[i];
                if (p.error_model_active) {
                    const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
                    const double bi = p.error_model_b * I_for_b;
                    const double v = p.error_model_a * double(sigma_corr) * sigma_corr + bi * bi;
                    if (v > 0.0) sigma_corr = float(sqrt(v));
                }
                if (p.reject_outliers && p.error_model_active && isfinite(rmed)
                    && fabsf(I_corr - rmed) > p.reject_nsigma * sigma_corr) {
                    ++rejected; p.rejected_obs[i] = 1; continue;
                }
                const double w = 1.0 / (double(sigma_corr) * sigma_corr);
                const double wI = w * I_corr;
                const int half = HalfForImageDev(p.frame[i]);
                swI += wI; sw += w;
                if (half) { swIh1 += wI; swh1 += w; ++nh1; } else { swIh0 += wI; swh0 += w; ++nh0; }
                if (!isfinite(dd) && isfinite(p.d[i]) && p.d[i] > 0.0f) dd = p.d[i];
            }
            p.a_swI[g] = swI; p.a_sw[g] = sw; p.a_swIh0[g] = swIh0; p.a_swIh1[g] = swIh1;
            p.a_swh0[g] = swh0; p.a_swh1[g] = swh1; p.a_nh0[g] = nh0; p.a_nh1[g] = nh1; p.a_d[g] = dd;
            p.a_rejected[g] = rejected;
        }
    }

    // One thread per group: R_meas accumulators (sum|I_corr - merged_I|, sum_I, n) + usable count for the
    // per-shell total_observations. Mirrors MergeAndStats' R_meas re-walk (looser, cell-only filter).
    __global__ void MergeRmeasKernel(MergeParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            const double mI = p.merged_I[g];
            const bool have = isfinite(mI);
            double absdev = 0, sumI = 0;
            int n = 0, nusable = 0;
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!RmeasUsable(i, p)) continue;
                if (p.rejected_obs[i]) continue;   // outlier-rejected in the merge -> also out of R_meas (XDS convention)
                if (!isfinite(p.d[i]) || !(p.d[i] > 0.0f)) continue;   // host counts only fulls with a shell
                ++nusable;
                if (have) {
                    const double I_corr = double(p.I[i]) * p.corr[i];
                    absdev += fabs(I_corr - mI); sumI += I_corr; ++n;
                }
            }
            p.r_absdev[g] = absdev; p.r_sumI[g] = sumI; p.r_n[g] = n; p.r_nusable[g] = nusable;
        }
    }

    void CudaCheck(cudaError_t e, const char *what) {
        if (e != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e));
    }

    template <typename T>
    void Upload(CudaDevicePtr<T> &dst, const T *src, int n) {
        dst = CudaDevicePtr<T>(std::max(1, n));
        if (n > 0)
            CudaCheck(cudaMemcpy(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice), "upload");
    }
}

struct RotationScaleMergeGPU::Impl {
    bool available = false;
    int n_obs = 0, n_frames = 0, n_groups = 0;

    // immutable per-obs
    CudaDevicePtr<float> I, sigma, rlp, partiality, zeta;
    CudaDevicePtr<uint8_t> on_ice;
    CudaDevicePtr<int32_t> frame;
    CudaDevicePtr<float> corr;               // mutable, resident across iterations
    CudaDevicePtr<int32_t> frame_start, frame_count;
    // per space group
    CudaDevicePtr<int32_t> group, group_perm, group_start, group_count;
    // scratch
    CudaDevicePtr<double> group_mean, g;
    CudaDevicePtr<uint8_t> scaled;
    CudaDevicePtr<float> sco_coeff;
    CudaDevicePtr<uint8_t> sco_ok;
    CudaDevicePtr<double> cc;        // per-frame CC (diagnostic), length n_frames
    CudaDevicePtr<int64_t> cc_n;
    CudaDevicePtr<uint8_t> smooth_apply;   // per-frame smooth-G apply flag + ratio, length n_frames
    CudaDevicePtr<double> smooth_ratio;
    // merge / error-model reductions over the resident fulls (reuse the fulls group CSR f_gperm/...)
    CudaDevicePtr<uint8_t> frame_cell_ok, merge_masked;
    CudaDevicePtr<double> m_sw, m_swI, m_em_mean;   // per group (n_groups)
    CudaDevicePtr<int32_t> m_cnt;
    CudaDevicePtr<double> m_s2, m_I2, m_dev2;        // per full (n_fulls)
    CudaDevicePtr<uint8_t> m_valid;
    CudaDevicePtr<uint8_t> m_rejected;   // per-full outlier-rejected flag (MergeAccum -> MergeRmeas)
    CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_d;   // merge accum per group
    CudaDevicePtr<int32_t> a_nh0, a_nh1, a_rejected;
    CudaDevicePtr<float> reject_median;
    CudaDevicePtr<double> merged_I, r_absdev, r_sumI;   // R_meas per group (merged_I uploaded)
    CudaDevicePtr<int32_t> r_n, r_nusable;
    int merge_for_search = 0, merge_n_masked = 0;   // filter context for one MergeAndStats call
    float merge_ice_hw = 0.0f;
    double merge_min_part = 0.0;

    // combine: extra per-obs inputs + the one-time raw-hkl run layout
    CudaDevicePtr<float> bkg, image_number, d_obs, px_obs, py_obs;
    int n_runs = 0, n_perm = 0;
    CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l, rr_group;
    CudaDevicePtr<int32_t> rr_nevents, rr_nusable, rr_offset;
    // combine: resident fulls SoA (rebuilt each Combine)
    int n_fulls = 0;
    CudaDevicePtr<int32_t> f_h, f_k, f_l, f_frame, f_group;
    CudaDevicePtr<float> f_I, f_sigma, f_d, f_img, f_px, f_py;
    CudaDevicePtr<uint8_t> f_on_ice;
    // scale-fulls (Unity model, kept resident): all-ones partiality/rlp/zeta so the shared scaling kernels
    // yield coeff=mean, plus the working corr, the per-obs scale scratch, and the fulls frame/group CSRs
    // (built on the host from the small f_frame/f_group key arrays, over the emit-ordered fulls).
    CudaDevicePtr<float> f_corr, f_partiality, f_rlp, f_zeta, f_sco_coeff;
    CudaDevicePtr<uint8_t> f_sco_ok;
    CudaDevicePtr<int32_t> f_frame_perm, f_frame_start, f_frame_count;
    CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
};

RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique<Impl>()) {
    if (get_gpu_count() > 0) {
        set_gpu(0);
        impl_->available = true;
    }
}

RotationScaleMergeGPU::~RotationScaleMergeGPU() = default;

bool RotationScaleMergeGPU::Available() const { return impl_->available; }

void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames,
                                        const float *I, const float *sigma, const float *rlp,
                                        const float *partiality, const float *zeta, const uint8_t *on_ice,
                                        const int32_t *frame, const float *corr0,
                                        const int32_t *frame_start, const int32_t *frame_count) {
    auto &d = *impl_;
    d.n_obs = n_obs;
    d.n_frames = n_frames;
    Upload(d.I, I, n_obs); Upload(d.sigma, sigma, n_obs); Upload(d.rlp, rlp, n_obs);
    Upload(d.partiality, partiality, n_obs); Upload(d.zeta, zeta, n_obs); Upload(d.on_ice, on_ice, n_obs);
    Upload(d.frame, frame, n_obs); Upload(d.corr, corr0, n_obs);
    Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames);
    d.g = CudaDevicePtr<double>(n_frames);
    d.scaled = CudaDevicePtr<uint8_t>(n_frames);
    d.sco_coeff = CudaDevicePtr<float>(n_obs);
    d.sco_ok = CudaDevicePtr<uint8_t>(n_obs);
    d.cc = CudaDevicePtr<double>(std::max(1, n_frames));
    d.cc_n = CudaDevicePtr<int64_t>(std::max(1, n_frames));
}

void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const int32_t *group_perm,
                                      int n_group_perm, const int32_t *group_start,
                                      const int32_t *group_count) {
    auto &d = *impl_;
    d.n_groups = n_groups;
    Upload(d.group, group, d.n_obs);
    Upload(d.group_perm, group_perm, n_group_perm);   // obs with group >= 0, in group order
    Upload(d.group_start, group_start, n_groups);
    Upload(d.group_count, group_count, n_groups);
    d.group_mean = CudaDevicePtr<double>(std::max(1, n_groups));
}

void RotationScaleMergeGPU::SetCorr(const float *corr) {
    CudaCheck(cudaMemcpy(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float),
                         cudaMemcpyHostToDevice), "upload corr");
}

void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min_partiality,
                                          bool /*has_d_min*/) {
    auto &d = *impl_;
    CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset scaled");
    CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset g"); // unscaled g unused
    const int obs_blocks = (d.n_obs + BLK - 1) / BLK;
    const int upd_blocks = std::min(65535, obs_blocks);
    const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
    for (int it = 0; it < iters; ++it) {
        ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
            d.group_perm.get(), d.group_start.get(), d.group_count.get(),
            d.I.get(), d.sigma.get(), d.partiality.get(), d.corr.get(), d.group_mean.get());
        PrepScaleObsKernel<<<obs_blocks, BLK>>>(d.n_obs, d.group.get(), d.partiality.get(), d.rlp.get(),
            d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sco_coeff.get(), d.sco_ok.get());
        FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k, d.frame_start.get(), d.frame_count.get(),
            d.I.get(), d.sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, d.g.get(), d.scaled.get());
        UpdateCorrKernel<<<upd_blocks, BLK>>>(d.n_obs, d.frame.get(), d.rlp.get(), d.partiality.get(),
            d.g.get(), d.scaled.get(), d.corr.get());
    }
    CudaCheck(cudaGetLastError(), "kernel launch");
    CudaCheck(cudaDeviceSynchronize(), "scale sync");
}

void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
    CudaCheck(cudaMemcpy(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float),
                         cudaMemcpyDeviceToHost), "download corr");
}

void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
    CudaCheck(cudaMemcpy(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double),
                         cudaMemcpyDeviceToHost), "download g");
    CudaCheck(cudaMemcpy(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t),
                         cudaMemcpyDeviceToHost), "download scaled");
}

void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
    Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames);
}

// The per-group inv-var mean (em_mean) + the per-full leverage-corrected error-model samples over the
// resident+scaled fulls. Stashes the filter context for the later MergeAccum/MergeRmeas calls.
void RotationScaleMergeGPU::MergeEmSamples(bool for_search, const uint8_t *masked, int n_masked,
                                           double ice_half_width_q, double min_partiality,
                                           double *em_mean_out, int32_t *cnt_out, double *s2_out,
                                           double *I2_out, double *dev2_out, uint8_t *valid_out) {
    auto &d = *impl_;
    const int ng = d.n_groups, nf = d.n_fulls;
    d.merge_for_search = for_search ? 1 : 0; d.merge_n_masked = n_masked;
    d.merge_ice_hw = float(ice_half_width_q); d.merge_min_part = min_partiality;
    d.m_sw = CudaDevicePtr<double>(std::max(1, ng)); d.m_swI = CudaDevicePtr<double>(std::max(1, ng));
    d.m_em_mean = CudaDevicePtr<double>(std::max(1, ng)); d.m_cnt = CudaDevicePtr<int32_t>(std::max(1, ng));
    d.m_s2 = CudaDevicePtr<double>(std::max(1, nf)); d.m_I2 = CudaDevicePtr<double>(std::max(1, nf));
    d.m_dev2 = CudaDevicePtr<double>(std::max(1, nf)); d.m_valid = CudaDevicePtr<uint8_t>(std::max(1, nf));
    Upload(d.merge_masked, masked, n_masked);

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = min_partiality; p.ice_hw = d.merge_ice_hw;
    p.for_search = d.merge_for_search; p.n_masked = n_masked;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get(); p.masked = d.merge_masked.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.em_mean = d.m_em_mean.get(); p.sw = d.m_sw.get(); p.swI = d.m_swI.get();
    p.em_mean_out = d.m_em_mean.get(); p.cnt = d.m_cnt.get();
    p.s2 = d.m_s2.get(); p.I2 = d.m_I2.get(); p.dev2 = d.m_dev2.get(); p.valid = d.m_valid.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
    MergeEmStatsKernel<<<grp_blocks, BLK>>>(p);
    MergeSamplesKernel<<<obs_blocks, BLK>>>(nf, p);
    CudaCheck(cudaGetLastError(), "merge em/samples launch");
    CudaCheck(cudaDeviceSynchronize(), "merge em/samples sync");
    CudaCheck(cudaMemcpy(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double),
                         cudaMemcpyDeviceToHost), "dl em_mean");
    CudaCheck(cudaMemcpy(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t),
                         cudaMemcpyDeviceToHost), "dl cnt");
    if (nf > 0) {
        CudaCheck(cudaMemcpy(s2_out, d.m_s2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl s2");
        CudaCheck(cudaMemcpy(I2_out, d.m_I2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl I2");
        CudaCheck(cudaMemcpy(dev2_out, d.m_dev2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl dev2");
        CudaCheck(cudaMemcpy(valid_out, d.m_valid.get(), size_t(nf) * sizeof(uint8_t), cudaMemcpyDeviceToHost), "dl valid");
    }
}

void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
                                       bool reject_outliers, double reject_nsigma, const float *reject_median,
                                       double *swI, double *sw, double *swIh0, double *swIh1,
                                       double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out,
                                       int32_t *rejected) {
    auto &d = *impl_;
    const int ng = d.n_groups;
    d.a_swI = CudaDevicePtr<double>(std::max(1, ng)); d.a_sw = CudaDevicePtr<double>(std::max(1, ng));
    d.a_swIh0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swIh1 = CudaDevicePtr<double>(std::max(1, ng));
    d.a_swh0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swh1 = CudaDevicePtr<double>(std::max(1, ng));
    d.a_nh0 = CudaDevicePtr<int32_t>(std::max(1, ng)); d.a_nh1 = CudaDevicePtr<int32_t>(std::max(1, ng));
    d.a_d = CudaDevicePtr<double>(std::max(1, ng)); d.a_rejected = CudaDevicePtr<int32_t>(std::max(1, ng));
    d.m_rejected = CudaDevicePtr<uint8_t>(std::max(1, d.n_fulls));
    CudaCheck(cudaMemset(d.m_rejected.get(), 0, size_t(std::max(1, d.n_fulls)) * sizeof(uint8_t)),
              "zero m_rejected");
    Upload(d.reject_median, reject_median, ng);

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part; p.ice_hw = d.merge_ice_hw;
    p.for_search = d.merge_for_search; p.n_masked = d.merge_n_masked;
    p.error_model_a = error_model_a; p.error_model_b = error_model_b;
    p.error_model_active = error_model_active ? 1 : 0;
    p.reject_outliers = reject_outliers ? 1 : 0; p.reject_nsigma = reject_nsigma;
    p.reject_median = d.reject_median.get();
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get(); p.masked = d.merge_masked.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.em_mean = d.m_em_mean.get();
    p.a_swI = d.a_swI.get(); p.a_sw = d.a_sw.get(); p.a_swIh0 = d.a_swIh0.get(); p.a_swIh1 = d.a_swIh1.get();
    p.a_swh0 = d.a_swh0.get(); p.a_swh1 = d.a_swh1.get(); p.a_nh0 = d.a_nh0.get(); p.a_nh1 = d.a_nh1.get();
    p.a_d = d.a_d.get(); p.a_rejected = d.a_rejected.get();
    p.rejected_obs = d.m_rejected.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    MergeAccumKernel<<<grp_blocks, BLK>>>(p);
    CudaCheck(cudaGetLastError(), "merge accum launch");
    CudaCheck(cudaDeviceSynchronize(), "merge accum sync");
    auto dl = [&](void *h, const CudaDevicePtr<double> &s) {
        CudaCheck(cudaMemcpy(h, s.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl accum"); };
    dl(swI, d.a_swI); dl(sw, d.a_sw); dl(swIh0, d.a_swIh0); dl(swIh1, d.a_swIh1);
    dl(swh0, d.a_swh0); dl(swh1, d.a_swh1); dl(d_out, d.a_d);
    CudaCheck(cudaMemcpy(nh0, d.a_nh0.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh0");
    CudaCheck(cudaMemcpy(nh1, d.a_nh1.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh1");
    CudaCheck(cudaMemcpy(rejected, d.a_rejected.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rej");
}

void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI,
                                       int32_t *n, int32_t *nusable) {
    auto &d = *impl_;
    const int ng = d.n_groups;
    Upload(d.merged_I, merged_I, ng);
    d.r_absdev = CudaDevicePtr<double>(std::max(1, ng)); d.r_sumI = CudaDevicePtr<double>(std::max(1, ng));
    d.r_n = CudaDevicePtr<int32_t>(std::max(1, ng)); d.r_nusable = CudaDevicePtr<int32_t>(std::max(1, ng));

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.merged_I = d.merged_I.get();
    p.r_absdev = d.r_absdev.get(); p.r_sumI = d.r_sumI.get(); p.r_n = d.r_n.get(); p.r_nusable = d.r_nusable.get();
    p.rejected_obs = d.m_rejected.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    MergeRmeasKernel<<<grp_blocks, BLK>>>(p);
    CudaCheck(cudaGetLastError(), "merge rmeas launch");
    CudaCheck(cudaDeviceSynchronize(), "merge rmeas sync");
    CudaCheck(cudaMemcpy(absdev, d.r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl absdev");
    CudaCheck(cudaMemcpy(sumI, d.r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl sumI");
    CudaCheck(cudaMemcpy(n, d.r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rn");
    CudaCheck(cudaMemcpy(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rnusable");
}

void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) {
    auto &d = *impl_;
    Upload(d.smooth_apply, apply, d.n_frames);
    Upload(d.smooth_ratio, ratio, d.n_frames);
    const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
    SmoothCorrKernel<<<blocks, BLK>>>(d.n_obs, d.frame.get(), d.smooth_apply.get(),
                                      d.smooth_ratio.get(), d.corr.get());
    CudaCheck(cudaGetLastError(), "smooth corr launch");
    CudaCheck(cudaDeviceSynchronize(), "smooth corr sync");
}

void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out) {
    auto &d = *impl_;
    const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
    // Post-smooth group means (reuse the scaling reduce; reads the resident, smoothed corr), then the
    // per-frame CC over the resident partials. Only the tiny per-frame cc/cc_n come back to the host.
    ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
        d.group_perm.get(), d.group_start.get(), d.group_count.get(),
        d.I.get(), d.sigma.get(), d.partiality.get(), d.corr.get(), d.group_mean.get());
    PerFrameCCKernel<<<d.n_frames, BLK>>>(d.n_frames, min_partiality,
        d.frame_start.get(), d.frame_count.get(), d.I.get(), d.sigma.get(), d.partiality.get(),
        d.corr.get(), d.on_ice.get(), d.group.get(), d.group_mean.get(), d.cc.get(), d.cc_n.get());
    CudaCheck(cudaGetLastError(), "partial CC launch");
    CudaCheck(cudaDeviceSynchronize(), "partial CC sync");
    CudaCheck(cudaMemcpy(cc_out, d.cc.get(), size_t(d.n_frames) * sizeof(double),
                         cudaMemcpyDeviceToHost), "download cc");
    CudaCheck(cudaMemcpy(cc_n_out, d.cc_n.get(), size_t(d.n_frames) * sizeof(int64_t),
                         cudaMemcpyDeviceToHost), "download cc_n");
}

void RotationScaleMergeGPU::SetCombineInputs(const float *bkg, const float *image_number, const float *d,
                                            const float *px, const float *py) {
    auto &dd = *impl_;
    Upload(dd.bkg, bkg, dd.n_obs);
    Upload(dd.image_number, image_number, dd.n_obs);
    Upload(dd.d_obs, d, dd.n_obs);
    Upload(dd.px_obs, px, dd.n_obs);
    Upload(dd.py_obs, py, dd.n_obs);
}

void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *perm,
                                       const int32_t *rr_start, const int32_t *rr_count,
                                       const int32_t *rr_h, const int32_t *rr_k, const int32_t *rr_l) {
    auto &d = *impl_;
    d.n_runs = n_runs;
    d.n_perm = n_perm;
    Upload(d.perm, perm, n_perm);
    Upload(d.rr_start, rr_start, n_runs);
    Upload(d.rr_count, rr_count, n_runs);
    Upload(d.rr_h, rr_h, n_runs);
    Upload(d.rr_k, rr_k, n_runs);
    Upload(d.rr_l, rr_l, n_runs);
    d.rr_group = CudaDevicePtr<int32_t>(std::max(1, n_runs));
    d.rr_nevents = CudaDevicePtr<int32_t>(std::max(1, n_runs));
    d.rr_nusable = CudaDevicePtr<int32_t>(std::max(1, n_runs));
    d.rr_offset = CudaDevicePtr<int32_t>(std::max(1, n_runs));
}

int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
                                   double capture_uncertainty_coeff, double min_captured_fraction) {
    auto &d = *impl_;
    CudaCheck(cudaMemcpy(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
                         cudaMemcpyHostToDevice), "upload rr_group");

    CombineParams p{};
    p.n_runs = d.n_runs;
    p.min_partiality = min_partiality;
    p.capture_uncertainty_coeff = capture_uncertainty_coeff;
    p.min_captured_fraction = min_captured_fraction;
    p.I = d.I.get(); p.sigma = d.sigma.get(); p.corr = d.corr.get(); p.partiality = d.partiality.get();
    p.bkg = d.bkg.get(); p.image_number = d.image_number.get(); p.d = d.d_obs.get();
    p.px = d.px_obs.get(); p.py = d.py_obs.get();
    p.frame = d.frame.get(); p.on_ice = d.on_ice.get();
    p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get();
    p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = d.rr_group.get();
    p.rr_nevents = d.rr_nevents.get(); p.rr_nusable = d.rr_nusable.get();

    const int blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);

    // Count pass: how many fulls each run emits.
    CombineKernel<false><<<blocks, BLK>>>(p);
    CudaCheck(cudaGetLastError(), "combine count launch");

    // Exclusive prefix sum on the host (deterministic) -> per-run output offset + total fulls.
    std::vector<int32_t> nevents(d.n_runs);
    CudaCheck(cudaMemcpy(nevents.data(), d.rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
                         cudaMemcpyDeviceToHost), "download nevents");
    std::vector<int32_t> offset(d.n_runs);
    int64_t acc = 0;
    for (int r = 0; r < d.n_runs; ++r) { offset[r] = static_cast<int32_t>(acc); acc += nevents[r]; }
    d.n_fulls = static_cast<int>(acc);

    // Allocate the fulls SoA and emit.
    const int nf = std::max(1, d.n_fulls);
    d.f_h = CudaDevicePtr<int32_t>(nf); d.f_k = CudaDevicePtr<int32_t>(nf); d.f_l = CudaDevicePtr<int32_t>(nf);
    d.f_frame = CudaDevicePtr<int32_t>(nf); d.f_group = CudaDevicePtr<int32_t>(nf);
    d.f_I = CudaDevicePtr<float>(nf); d.f_sigma = CudaDevicePtr<float>(nf);
    d.f_d = CudaDevicePtr<float>(nf); d.f_img = CudaDevicePtr<float>(nf);
    d.f_px = CudaDevicePtr<float>(nf); d.f_py = CudaDevicePtr<float>(nf);
    d.f_on_ice = CudaDevicePtr<uint8_t>(nf);
    d.f_corr = CudaDevicePtr<float>(nf); d.f_partiality = CudaDevicePtr<float>(nf);
    d.f_rlp = CudaDevicePtr<float>(nf); d.f_zeta = CudaDevicePtr<float>(nf);
    d.f_sco_coeff = CudaDevicePtr<float>(nf); d.f_sco_ok = CudaDevicePtr<uint8_t>(nf);
    CudaCheck(cudaMemcpy(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t),
                         cudaMemcpyHostToDevice), "upload offset");

    p.rr_offset = d.rr_offset.get();
    p.f_h = d.f_h.get(); p.f_k = d.f_k.get(); p.f_l = d.f_l.get();
    p.f_frame = d.f_frame.get(); p.f_group = d.f_group.get();
    p.f_I = d.f_I.get(); p.f_sigma = d.f_sigma.get(); p.f_d = d.f_d.get(); p.f_img = d.f_img.get();
    p.f_px = d.f_px.get(); p.f_py = d.f_py.get();
    p.f_on_ice = d.f_on_ice.get();
    if (d.n_fulls > 0) {
        CombineKernel<true><<<blocks, BLK>>>(p);
        CudaCheck(cudaGetLastError(), "combine emit launch");
    }
    CudaCheck(cudaDeviceSynchronize(), "combine sync");
    return d.n_fulls;
}

void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float *I, float *sigma, float *d,
                                     float *image_number, int32_t *frame, uint8_t *on_ice,
                                     int32_t *group) const {
    const auto &dd = *impl_;
    const size_t n = static_cast<size_t>(dd.n_fulls);
    if (n == 0) return;
    auto dl = [&](void *dst, const void *src, size_t bytes) {
        CudaCheck(cudaMemcpy(dst, src, bytes, cudaMemcpyDeviceToHost), "download fulls");
    };
    dl(h, dd.f_h.get(), n * sizeof(int32_t)); dl(k, dd.f_k.get(), n * sizeof(int32_t));
    dl(l, dd.f_l.get(), n * sizeof(int32_t)); dl(frame, dd.f_frame.get(), n * sizeof(int32_t));
    dl(group, dd.f_group.get(), n * sizeof(int32_t));
    dl(I, dd.f_I.get(), n * sizeof(float)); dl(sigma, dd.f_sigma.get(), n * sizeof(float));
    dl(d, dd.f_d.get(), n * sizeof(float)); dl(image_number, dd.f_img.get(), n * sizeof(float));
    dl(on_ice, dd.f_on_ice.get(), n * sizeof(uint8_t));
}

void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) const {
    const auto &d = *impl_;
    if (d.n_fulls == 0) return;
    const size_t bytes = size_t(d.n_fulls) * sizeof(int32_t);
    CudaCheck(cudaMemcpy(frame, d.f_frame.get(), bytes, cudaMemcpyDeviceToHost), "download f_frame");
    CudaCheck(cudaMemcpy(group, d.f_group.get(), bytes, cudaMemcpyDeviceToHost), "download f_group");
}

void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_perm,
                                             const int32_t *frame_start, const int32_t *frame_count) {
    auto &d = *impl_;
    Upload(d.f_frame_perm, frame_perm, n_perm);
    Upload(d.f_frame_start, frame_start, d.n_frames);
    Upload(d.f_frame_count, frame_count, d.n_frames);
}

void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm,
                                           const int32_t *gstart, const int32_t *gcount) {
    auto &d = *impl_;
    Upload(d.f_gperm, gperm, n_gperm);
    Upload(d.f_gstart, gstart, d.n_groups);
    Upload(d.f_gcount, gcount, d.n_groups);
}

void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_partiality) {
    auto &d = *impl_;
    const int nf = d.n_fulls;
    if (nf == 0) return;
    const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
    const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);

    // Unity model: partiality/rlp/zeta = 1 so coeff = mean; corr starts at 1.
    FillKernel<<<obs_blocks, BLK>>>(d.f_corr.get(), nf, 1.0f);
    FillKernel<<<obs_blocks, BLK>>>(d.f_partiality.get(), nf, 1.0f);
    FillKernel<<<obs_blocks, BLK>>>(d.f_rlp.get(), nf, 1.0f);
    FillKernel<<<obs_blocks, BLK>>>(d.f_zeta.get(), nf, 1.0f);
    // g/scaled reset once (a frame, once fitted, stays fitted - same as ScalePartials).
    CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset f scaled");
    CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset f g");

    for (int it = 0; it < iters; ++it) {
        ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
            d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(),
            d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get());
        PrepScaleObsKernel<<<obs_blocks, BLK>>>(nf, d.f_group.get(), d.f_partiality.get(), d.f_rlp.get(),
            d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), d.f_sco_coeff.get(), d.f_sco_ok.get());
        FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k,
            d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_sigma.get(),
            d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), d.g.get(), d.scaled.get());
        UpdateCorrKernel<<<obs_blocks, BLK>>>(nf, d.f_frame.get(), d.f_rlp.get(), d.f_partiality.get(),
            d.g.get(), d.scaled.get(), d.f_corr.get());
    }
    CudaCheck(cudaGetLastError(), "scale fulls launch");
    CudaCheck(cudaDeviceSynchronize(), "scale fulls sync");
}

void RotationScaleMergeGPU::GetFullsCorr(float *corr) const {
    const auto &d = *impl_;
    if (d.n_fulls == 0) return;
    CudaCheck(cudaMemcpy(corr, d.f_corr.get(), size_t(d.n_fulls) * sizeof(float),
                         cudaMemcpyDeviceToHost), "download f_corr");
}

void RotationScaleMergeGPU::GetFullsPxPy(float *px, float *py) const {
    const auto &d = *impl_;
    if (d.n_fulls == 0) return;
    const size_t bytes = size_t(d.n_fulls) * sizeof(float);
    CudaCheck(cudaMemcpy(px, d.f_px.get(), bytes, cudaMemcpyDeviceToHost), "download f_px");
    CudaCheck(cudaMemcpy(py, d.f_py.get(), bytes, cudaMemcpyDeviceToHost), "download f_py");
}

void RotationScaleMergeGPU::SetFullsCorr(const float *corr) {
    auto &d = *impl_;
    if (d.n_fulls == 0) return;
    CudaCheck(cudaMemcpy(d.f_corr.get(), corr, size_t(d.n_fulls) * sizeof(float),
                         cudaMemcpyHostToDevice), "upload f_corr");
}