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

#pragma once

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

#include "ErrorModel.h"

struct Reflection;

// GPU engine for the RotationScaleMerge hot loops. The class keeps the per-observation data resident on
// the device as a structure-of-arrays (coalesced) and runs the scaling loop there. The host keeps the
// one-time raw-hkl sort and the per-space-group ASU keying (gemmi); it hands the GPU the dense group ids
// plus a group-ordered permutation so the per-group reduction is a deterministic segmented reduction
// (one block per group, fixed order, no atomics) - matching the run-to-run determinism of the CPU path.
//
// Only compiled when CUDA is available; the header is safe to include unconditionally (the impl behind
// the pimpl is null without CUDA, and RotationScaleMerge falls back to the CPU loops).
class RotationScaleMergeGPU {
public:
    // On `device`, which must be one of the visible GPUs.
    explicit RotationScaleMergeGPU(int device = 0);
    ~RotationScaleMergeGPU();
    RotationScaleMergeGPU(const RotationScaleMergeGPU &) = delete;
    RotationScaleMergeGPU &operator=(const RotationScaleMergeGPU &) = delete;

    // True if a GPU was found and the engine is usable.
    [[nodiscard]] bool Available() const;

    // The observation count and the frame CSR (frame_start/frame_count, length n_frames, indexing
    // the obs arrays in frame order), plus every device-side observation buffer, allocated here so
    // the chunked uploads below have somewhere to land. Call once, before them.
    void SetPartialsLayout(int n_obs, int n_frames,
                           const int32_t *frame_start, const int32_t *frame_count);

    // One slice ([offset, offset+count) of n_obs) of one immutable per-observation field, so the
    // host can stage a bounded chunk at a time out of its AoS record: staged whole and all at once,
    // the arrays are more than half the entire observation payload held live again.
    // Corr0 seeds the resident, mutable corr (SetCorr refreshes it). Bkg/VarBkg/ImageNumber/D
    // are the combine inputs (image-local background, fractional frame position for event
    // contiguity and resolution; the predicted detector position carried to the full for the
    // absorption surface stays on the host, GetFullsPeak).
    enum class ObsField { I, Sigma, PrescalingCorr, Partiality, Zeta, Corr0,
                          Bkg, VarBkg, ImageNumber, D };
    void SetObsField(ObsField f, int offset, int count, const float *v);
    // The flags byte: on an ice ring (bit 0), then the OBS_CLIPPED / OBS_OVERLOADED bits of
    // RotationScaleMerge one bit up. The frame is not uploaded: the CSR above gives it.
    void SetObsFlags(int offset, int count, const uint8_t *flags);

    // Every per-observation field the uploads above set, the frame CSR and the raw-hkl runs, copied from
    // `from` (which has them all, on its own device or this one) instead of being staged again from the
    // host. Call once, in place of SetPartialsLayout, the uploads and SetRawRuns, before anything changes
    // the fields on `from`.
    void CopyObservationsFrom(const RotationScaleMergeGPU &from);

    // --- the resident ingest (RotationScaleMerge::IngestDevice with no observation array on the host) ---
    // The steps below build, from the reflections, exactly the observations and raw-hkl runs that
    // RotationScaleMerge::BuildInRangeObservations builds on the host, and leave them where the uploads
    // above and SetRawRuns would: the reflections are sorted and cut on the card, and cross the bus
    // twice - once for what the sort and the cut read, once for the observations the cut kept.

    // The reflections of n_frames frames, frame o's count[o] of them at refl[o], uploaded in frame order;
    // what the sort, the cut and the background below read is kept per reflection.
    void IngestUpload(int n_frames, const Reflection *const *refl, const int32_t *count, size_t nthreads);
    // Each frame's mean background over its reflections with a finite, positive one, summed in their
    // order (NaN where there is none): the incident-flux meter. Length n_frames.
    void IngestMeanBackground(double *mean_bkg);
    // Order the reflections by (raw h,k,l, image_number, index) and cut that order into raw-hkl runs,
    // each with the resolution of its first member that has a usable one; and drop the runs outside
    // [d_min, d_max] (either may be absent) with their reflections. Returns how many observations are
    // left. Nothing per observation is on the card yet (IngestObservations).
    int IngestSort(std::optional<double> d_min, std::optional<double> d_max);
    // The layout IngestSort left: the reflection and run counts before the resolution cut, the
    // observations' frame CSR, and the runs that survived it.
    struct IngestLayout {
        int64_t n_reflections = 0;
        int n_runs_all = 0;
        std::vector<int32_t> frame_start, frame_count;
        std::vector<int32_t> run_start, run_count, run_h, run_k, run_l;
        std::vector<float> run_d;
    };
    void IngestGetLayout(IngestLayout &out) const;
    // The observations IngestSort kept, taken from the same reflections (as IngestUpload had them), in
    // its order, each with its deterministic correction divided by its frame's flux (length n_frames).
    // px/py (length IngestSort's count) receive each observation's predicted position, which is not
    // kept on the card (GetFullsPeak).
    void IngestObservations(const Reflection *const *refl, const int32_t *count, size_t nthreads,
                            const double *frame_flux, float *px, float *py);
    // The ingest's geometry smoothing, on the observations IngestSort left (RotationScaleMerge::IngestOnDevice):
    // each delta_phi re-solved against its frame's smoothed lattice where the frame has one - recip holds
    // a*, b*, c* per frame (9 floats), ok which frames, and the goniometer frame as SmoothedGeometry
    // has it - returning how many were; one exact-Bragg angle laid over each rocking event (counts of
    // events and of partials relaid); and the partiality recomputed from the per-frame mosaicity, with
    // corr brought to it. The host's loops, with every rounding explicit.
    struct DeltaPhiFrame { std::vector<float> recip; std::vector<uint8_t> ok;
                           float S0[3], m1[3], m2[3], m3[3]; float m2_S0, m3_S0; };
    int64_t IngestDeltaPhi(const DeltaPhiFrame &g);
    void IngestRelayEvents(double increment_deg, float max_frame_gap, int64_t &n_events, int64_t &n_partials);
    // The histogram of the frames each rocking event spans (length n_frames + 2, the last bin taking
    // every longer one), over the observations usable as integrated: RockingEventFrames' count.
    void IngestRockingHistogram(float max_frame_gap, int32_t *hist);
    void IngestPartiality(const float *mosaicity, double wedge_deg, double bandwidth_sigma, double half_wavelength_A);
    // corr as the ingest leaves it (length IngestSort's count); the transient ingest arrays go.
    void IngestFinish(float *corr);

    // Per space group: the dense ASU-group id per obs, and a group-ordered permutation of the obs whose
    // group >= 0, with its CSR - so each group's observations are a contiguous, fixed-order segment for
    // the reduction. Built on the device from the raw-hkl runs uploaded by SetRawRuns: rawrun_group is
    // each run's group (-1 = none), and group g holds the runs sorted_run[group_first[g] ..
    // group_first[g + 1]). An observation is in its run's group when it passes the finiteness test
    // Ingest applies, and each group lists its observations in ascending index order.
    void BuildGroups(int n_groups, const int32_t *rawrun_group, int n_sorted, const int32_t *sorted_run,
                     const int32_t *group_first);

    // Re-upload the working corr (length n_obs) before a scaling pass (the host mutates it via smooth-G
    // between passes). SetObsField(Corr0) uploads the initial corr; this refreshes it.
    void SetCorr(const float *corr);

    // Run `iters` of {reduce group means -> per-frame weighted-LS G -> update corr} on the device,
    // in place on the resident corr. Rotation model (partiality folded via the stored partiality).
    void ScalePartials(int iters, double min_partiality, bool has_d_min);

    // Copy the updated corr back to the host (length n_obs), and the fitted per-frame G (length
    // n_frames, as double) plus the per-frame "was fitted" flag.
    void GetCorr(float *corr_out) const;
    void GetG(double *g_out, uint8_t *scaled_out) const;
    // The information each frame's last fit had on G, sum w^2 c^2 (length n_frames, 0 = not fitted).
    void GetInfo(double *info_out) const;

    // Apply the smooth-G correction to the resident corr: corr[i] *= ratio[frame[i]] for frames with
    // apply[f] (both length n_frames), matching CPU SmoothG. Keeps corr resident (no round-trip).
    void SmoothCorr(const uint8_t *apply, const double *ratio);

    // The two pass filters, on the resident corr (both mirror the host loops in Run and keep corr on the
    // device). By zeta: zero corr wherever the rocking geometry fails the de-novo search threshold,
    // returning how many observations that removed from the merge, i.e. how many had a finite, positive
    // corr. By frame: zero corr on the frames flagged in `reject` (length n_frames).
    int64_t FilterCorrByZeta(double min_zeta);
    void FilterCorrByFrame(const uint8_t *reject);

    // --- merge + error-model reductions over the resident, scaled fulls (reuse the fulls group CSR) ---

    // The per-frame cell-consistency mask (length n_frames) used by the merge filter. Uploaded once.
    void SetFrameCellOk(const uint8_t *frame_cell_ok);

    // Per-group inv-var mean (em_mean, length n_groups) + the leverage-corrected error-model samples in
    // fulls order, each with the counting variance at its reflection's expected intensity - the host's
    // samples to the bit (MergeAndStats' first two error-model loops).
    // Stashes (for_search, min_partiality) for the MergeAccum/MergeRmeas calls that follow.
    // hand / has_hands are the per-full Bijvoet hand and the per-group "this group pools two hands",
    // which the error-model samples are formed on; both null leaves the fit on the pooled group.
    void MergeEmSamples(bool for_search, double min_partiality,
                        const uint8_t *hand, const uint8_t *has_hands,
                        double *em_mean_out, int32_t *cnt_out, std::vector<ErrorModelSample> &samples);

    // What the pooled outlier test's widening by the Bijvoet variance is measured from (MergeAndStats'
    // reject_var_add block), with the error model (a, b) of the first fit: per group (length n_groups)
    // its d - its first usable full's - and the excess dI^2 - var(dI) of its two hands (NaN where it has
    // not both); and the d range of the fulls that entered a hand. The host's numbers to the bit.
    // Requires MergeEmSamples with the hands first.
    void MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d, double *excess,
                            float &d_lo, float &d_hi);

    // The sums the strong-reflection ISa asymptote is read from (MergeAndStats), per group (length
    // n_groups): sum and sum of squares of I*corr over the usable fulls, the sum of their counting
    // variances at the group's em_mean, and their count. The host's numbers to the bit. Requires
    // MergeEmSamples first.
    void MergeGroupScatter(double *sum, double *sum_sq, double *sum_var, int32_t *n);

    // Every full's ObservationKey and its CC1/2 half - the parity of the key's rank among its group's
    // usable fulls, 0 for an unusable one (length n_fulls each); the host's AssignHalvesOfMembers to
    // the bit, as it is integer work. Requires MergeEmSamples first (the filter context).
    void MergeHalves(uint64_t *obs_key, uint8_t *half);

    // Per-group merge accumulators (inv-var sums + deterministic half-sets, error-model-corrected sigma
    // from a/b). `half` is the per-full CC1/2 half-set (length n_fulls) assigned on the host, so this
    // kernel and the host merge loop cannot disagree about it.
    // The per-group results stay on the device for MergeAccumRange to download; reject_median is
    // uploaded (NAN where none). rejected_obs is the per-full flag (length n_fulls): on entry the
    // observations the host already rejected (the Wilson test), on return those plus the median test's;
    // the host needs it for the reductions it still does itself, above all the anomalous I(+)/I(-) split.
    // frame_cc_factor (length n_frames) is each frame's (G_ref/G)^2 floored at 1; swh_typ0/1 are the
    // half-set weights multiplied by it. Requires MergeEmSamples first (em_mean resident).
    // reject_var_add (n_groups) widens the pooled cut by the shell's own measured Bijvoet
    // variance; null leaves the plain n-sigma test.
    // reject_mult_sd is the multiplicative spread of the outlier band (OutlierBand.h).
    void MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
                    bool reject_outliers, double reject_nsigma, double reject_mult_sd, const float *reject_median,
                    const float *reject_var_add,
                    const uint8_t *half, const double *frame_cc_factor, uint8_t *rejected_obs);
    // Download groups [g0, g0 + n) of what MergeAccum left on the device; every output has length n.
    // rejected[g] counts the outliers dropped from the group.
    void MergeAccumRange(int g0, int n, double *swI, double *sw, double *swIh0, double *swIh1,
                         double *swh0, double *swh1, double *swh_typ0, double *swh_typ1,
                         int32_t *nh0, int32_t *nh1, double *d_out,
                         int32_t *rejected, uint8_t *on_ice_out);
    // Hand back what MergeAccum left - thirteen numbers per group - once every range is downloaded, so
    // that the sums the merge goes on to take do not stand beside them.
    void MergeAccumRelease();

    // Per-group R_meas accumulators: sum|I_corr-merged_I| and sum I, the same with each observation
    // weighted by its merge weight v = 1/sigma^2 under the error model of the last MergeAccum,
    // sum v, sum v^2, n, and the count this looser walk accepted - which the host uses only to skip
    // empty groups, NOT as the per-shell total_observations; merged_I is uploaded. All arrays length
    // n_groups.
    void MergeRmeas(const double *merged_I, double *absdev, double *sumI, double *wabsdev, double *wsumI,
                    double *sumv, double *sumv2, int32_t *n, int32_t *nusable);

    // Post-smooth per-frame diagnostic CC: recompute the group means from the resident (smoothed) corr
    // and the Pearson CC of each frame's I*corr vs its group mean, downloading only the per-frame cc /
    // cc_n (length n_frames). Mirrors ReduceGroupMeans(partials) + FinalizePerFrameScale's CC loop.
    void ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out);

    // --- 3D combine (partials -> fulls), all on the device ---

    // The one-time raw-hkl run layout (space-group-independent): the (raw h,k,l, image_number)-sorted
    // permutation of the obs, split into contiguous per-raw-hkl runs. Uploaded once in Ingest.
    void SetRawRuns(int n_runs, int n_perm, const int32_t *perm,
                    const int32_t *rawrun_start, const int32_t *rawrun_count,
                    const int32_t *rawrun_h, const int32_t *rawrun_k, const int32_t *rawrun_l);

    // Combine the resident partials (reading the current resident corr) into fulls on the device,
    // mirroring RotationScaleMerge::Combine: one thread per raw-hkl run splits its usable partials into
    // rocking events (frame gap <= max_frame_gap, the host's RockingEventFrameGap), pools background, seeds F, does 3 de-biased Poisson reweights and
    // adds the capture-uncertainty term. rawrun_group (length n_runs) is the current space group's ASU
    // id per raw hkl (it becomes the full's group). Deterministic: fulls are emitted in raw-run-major,
    // event order (a count pass -> host prefix sum -> emit-at-offset), matching the CPU path. Returns the
    // number of fulls (call GetFulls with buffers of that length); n_overloaded receives how many rocking
    // events were dropped for a saturated pixel (see the host combine). The previous fulls and the
    // partials' group CSR go here, so the partials' scaling (ScalePartials, ComputePartialCC) needs
    // BuildGroups again after it.
    int Combine(const int32_t *rawrun_group, double min_partiality, double capture_uncertainty_coeff,
                double min_captured_fraction, float max_frame_gap, int64_t &n_overloaded);

    // Download the combined fulls SoA (length = Combine()'s return); run is the raw-hkl run each full is
    // of. The working corr is downloaded separately by GetFullsCorr (it is only meaningful after the
    // fulls scaling; otherwise the caller sets it).
    void GetFulls(int32_t *run, float *I, float *sigma, float *d,
                  float *image_number, int32_t *frame, uint8_t *on_ice, uint8_t *clipped,
                  int32_t *group) const;

    // Download the observation each full's predicted detector position is that of (its peak partial),
    // for the host absorption surface; the positions themselves stay on the host. Length = n_fulls.
    void GetFullsPeak(int32_t *peak) const;

    // Download the fulls' variance model, var(I) = var_bkg + var_per_I * I + (capture * I)^2. Length = n_fulls.
    void GetFullsVariance(float *var_bkg, float *var_per_I, float *capture) const;

    // Hand back the arrays only the fulls' scaling and the downloads above read (the records, zeta, the
    // frame CSR, the group means, image number, px/py and the clipped flags), once the fulls are on the
    // host: the merge and the surfaces after it read none of them, and the next arm makes its own.
    void FullsDownloaded();

    // Re-upload the fulls' working corr (length n_fulls) after the host correction surfaces (decay /
    // absorption) mutate it, so the resident merge reads the corrected scale.
    void SetFullsCorr(const float *corr);

    // --- scale the resident fulls on the device (Unity model), no round-trip ---

    // The fulls' per-frame CSR (each frame's fulls, for the per-frame fit) and per-ASU-group CSR, built on
    // the device by stable sorts of the emit-ordered fulls, so that each frame and each group lists its
    // fulls in index order as a counting sort does. Hands back the fulls per frame (length n_frames) and
    // the group CSR for the host's own walks: gstart (n_groups + 1 entries, the last the total) and gperm.
    void BuildFullsCSRs(int32_t *frame_count, std::vector<int32_t> &gstart, std::vector<int32_t> &gperm);

    // Start the Unity scaling of the resident fulls: corr, partiality, prescaling and zeta all 1, no
    // frame fitted yet. Once before the FitFullsScale rounds.
    void ResetFullsScale();

    // One Unity fit of the resident fulls (reduce group means -> per-frame LS G and its information,
    // GetG / GetInfo), leaving the fulls' working corr as it is: the host smooths the fitted scales
    // and applies them with SmoothFullsCorr. Requires BuildFullsCSRs and
    // ResetFullsScale.
    void FitFullsScale(double min_partiality);

    // The fulls' counterpart of SmoothCorr: f_corr[i] *= ratio[f_frame[i]] where apply[f].
    void SmoothFullsCorr(const uint8_t *apply, const double *ratio);

    // Download the fulls' working corr (length = n_fulls), valid after the fulls scaling.
    void GetFullsCorr(float *corr) const;

    // --- correction-surface fit (RotationScaleMerge::ApplyCellSurface) ---
    // The two passes every round of the fit makes over its terms, on the device; the host keeps the
    // per-cell step, the gauge and the cross-validation. Both sums are formed in exactly the order the
    // host forms them, so the fitted surface is the host's to the last bit (see the kernels).

    // One observation as the surface fit sees it: the host's term, uploaded as it stands.
    struct SurfaceTerm { float I, sigma, corr, d; int32_t cell, group; };

    // The terms (in fulls order) with each one's frame parity, the ASU-group CSR over them (gperm lists
    // the terms of group g at [gstart[g], gstart[g+1]), in fulls order) and the cell count.
    void SurfaceSetTerms(int n_terms, const SurfaceTerm *term, const uint8_t *parity,
                         int n_groups, const int32_t *gperm, const int32_t *gstart, int ncell);

    // The same terms, group CSR and three subsets, built on the device from the resident fulls instead of
    // uploaded: cell (length n_fulls, < 0 = no cell) and corr (the host's current corr, length n_fulls)
    // are the only inputs from the host. A full gives a term where it has a group, a positive finite
    // corr, is off the ice rings and has a cell (ApplyCellSurface's filter); `block` is the host fit's
    // reduction block. The same arrays, in the same orders, as SurfaceSetTerms + SurfaceSetSubset are
    // handed by ApplyCellSurface's host build. Returns the number of terms. Requires Combine first.
    int SurfaceBuildTerms(const int32_t *cell, const float *corr, double min_partiality, int ncell, int n_groups,
                          int block);

    // The resolution gauge's inputs for the terms SurfaceBuildTerms left: each group's shell (that of its
    // last term in fulls order; length n_groups) and each (cell, shell) bin's statistical weight
    // sum 1/(sigma corr)^2 in fulls order (length ncell * nshell), over nshell equal-occupancy shells in
    // 1/d^2 - ApplyCellSurface's numbers to the bit.
    void SurfaceShellWeights(int n_terms, int nshell, int32_t *group_shell, double *shw);

    // One subset of the terms (0 = even frames, 1 = odd, 2 = all), cut into the host's n_blocks
    // reduction blocks and ordered within each block by cell, keeping term order inside a cell:
    // block b, cell c is perm[seg_start[b * ncell + c], seg_start[b * ncell + c + 1]).
    void SurfaceSetSubset(int subset, int n_blocks, const int32_t *perm, const int32_t *seg_start);

    // The per-group reference sums sw / swI over the terms of frame parity `parity` (< 0 = all) with
    // the surface A (length ncell) applied. They stay on the device for SurfaceFitSums;
    // SurfaceGetReference downloads them (length n_groups each).
    void SurfaceReference(int parity, const double *A);
    void SurfaceGetReference(double *sw, double *swI) const;

    // The fit's per-cell sums over one subset against the last SurfaceReference and its A:
    // cross = sum w Is Iref and ref2 = sum w Iref^2 (length ncell each).
    void SurfaceFitSums(int subset, double *cross, double *ref2);

    // Hand back the surface fits' buffers - several bytes per term, sized by the fulls - once the last
    // surface is fitted: the merge after the surfaces needs the card, and the next pass's surfaces
    // build their own (SurfaceBuildTerms).
    void SurfaceRelease();

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
};