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

#include <algorithm>
#include "../../common/JFJochMath.h"
#include "BraggPredictionRotGPU.h"
#include "RockingSlice.h"
#include "RotationColumnRange.h"
#include "../SensorAbsorption.h"

#ifdef JFJOCH_USE_CUDA
#include "../indexing/CUDAMemHelpers.h"
#include <cuda_runtime.h>
#include <cub/cub.cuh>
#include <cmath>

static inline void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

namespace {
     __host__ __device__ inline bool is_odd(int v) { return (v & 1) != 0; }

     __host__ __device__ inline void cross3(float ax, float ay, float az,
                                  float bx, float by, float bz,
                                  float &cx, float &cy, float &cz) {
        cx = ay * bz - az * by;
        cy = az * bx - ax * bz;
        cz = ax * by - ay * bx;
    }

     __host__ __device__ inline float dot3(float ax, float ay, float az,
                                 float bx, float by, float bz) {
        return ax * bx + ay * by + az * bz;
    }

     __host__ __device__ inline void normalize3(float &x, float &y, float &z) {
        float len = sqrtf(x * x + y * y + z * z);
        if (len < 1e-12f) { x = 0.0f; y = 0.0f; z = 0.0f; return; }
        float inv = 1.0f / len;
        x *= inv; y *= inv; z *= inv;
    }

    __device__ inline int compute_reflections_rot(const KernelConstsRot &C, int h, int k, int l, Reflection out[2]) {
        if (h == 0 && k == 0 && l == 0)
            return 0;

         switch (C.centering) {
             case 'I':
                 if (is_odd(h + k + l)) return false;
                 break;
             case 'A':
                 if (is_odd(k + l)) return false;
                 break;
             case 'B':
                 if (is_odd(h + l)) return false;
                 break;
             case 'C':
                 if (is_odd(h + k)) return false;
                 break;
             case 'F':
                 if (is_odd(h + k) || is_odd(h + l) || is_odd(k + l)) return false;
                 break;
             case 'R': {
                 int mod = (-h + k + l) % 3;
                 if (mod < 0) mod += 3;
                 if (mod != 0) return false;
                 break;
             }
             default:
                 break;
         }

        // p0 = A* h + B* k + C* l
        float p0x = C.Astar.x * h + C.Bstar.x * k + C.Cstar.x * l;
        float p0y = C.Astar.y * h + C.Bstar.y * k + C.Cstar.y * l;
        float p0z = C.Astar.z * h + C.Bstar.z * k + C.Cstar.z * l;

        float p0_sq = p0x * p0x + p0y * p0y + p0z * p0z;
        if (p0_sq <= 0.0f || p0_sq > C.one_over_dmax_sq)
            return 0;

        float p0_m1 = p0x * C.m1.x + p0y * C.m1.y + p0z * C.m1.z;
        float p0_m2 = p0x * C.m2.x + p0y * C.m2.y + p0z * C.m2.z;
        float p0_m3 = p0x * C.m3.x + p0y * C.m3.y + p0z * C.m3.z;

        float rho_sq = p0_sq - (p0_m2 * p0_m2);
        float p_m3 = (-p0_sq / 2.0f - p0_m2 * C.m2_S0) / C.m3_S0;
        float p_m2 = p0_m2;

        if (rho_sq < p_m3 * p_m3) return 0;
        if (p0_sq > 4.0f * dot3(C.S0.x, C.S0.y, C.S0.z, C.S0.x, C.S0.y, C.S0.z)) return 0;

        float p_m1_pos = sqrtf(rho_sq - p_m3 * p_m3);
        float p_m1_arr[2] = {p_m1_pos, -p_m1_pos};

        // Effective rocking width: mosaicity broadened by the bandwidth term (consistent with CPU),
        // dtheta = (dlambda/lambda) tan(theta_B) in quadrature with sigma_M, zeta-free.
        float mos_eff_rad = C.mos_angle_rad;
        if (C.bandwidth_sigma > 0.0f) {
            float sin_theta = C.half_wavelength_A * sqrtf(p0_sq);
            float dphi_bw = C.bandwidth_sigma * sin_theta / sqrtf(1.0f - sin_theta * sin_theta);
            mos_eff_rad = sqrtf(C.mos_angle_rad * C.mos_angle_rad + dphi_bw * dphi_bw);
        }

        int count = 0;
        for (int idx = 0; idx < 2; ++idx) {
            float p_m1 = p_m1_arr[idx];

            float cosphi = (p_m1 * p0_m1 + p_m3 * p0_m3) / rho_sq;
            float sinphi = (p_m1 * p0_m3 - p_m3 * p0_m1) / rho_sq;

            float px = C.m1.x * p_m1 + C.m2.x * p_m2 + C.m3.x * p_m3;
            float py = C.m1.y * p_m1 + C.m2.y * p_m2 + C.m3.y * p_m3;
            float pz = C.m1.z * p_m1 + C.m2.z * p_m2 + C.m3.z * p_m3;

            float Sx = C.S0.x + px;
            float Sy = C.S0.y + py;
            float Sz = C.S0.z + pz;

            float phi = -1.0f * atan2f(sinphi, cosphi);

            // e1 = normalize(S x S0) - direction perpendicular to both S and S0
            float e1x, e1y, e1z;
            cross3(Sx, Sy, Sz, C.S0.x, C.S0.y, C.S0.z, e1x, e1y, e1z);
            normalize3(e1x, e1y, e1z);

            // zeta = |m2 · e1| - the "lorentz-like" geometric factor for partiality
            float zeta_abs = fabsf(dot3(C.m2.x, C.m2.y, C.m2.z, e1x, e1y, e1z));

            // Check min_zeta threshold (consistent with CPU)
            if (zeta_abs < C.min_zeta)
                continue;

            // epsilon3 cutoff check (consistent with CPU, Kabsch formulation): measured against the
            // frame's EDGE, not its centre, or a reflection whose rocking curve overlaps the exposure
            // is rejected because its exact diffracting condition falls outside it - and since the
            // nearest frame centre is at most half a wedge away, it is then rejected on every frame
            // and lost entirely. See the CPU engine for the regime that reaches.
            float epsilon3 = fabsf(phi * zeta_abs) - 0.5f * C.wedge_angle_rad * zeta_abs;
            if (epsilon3 > C.mosaicity_multiplier * mos_eff_rad)
                continue;

            float cx, cy, cz;
            cross3(Sx, Sy, Sz, C.S0.x, C.S0.y, C.S0.z, cx, cy, cz);
            // Reciprocal Lorentz (Kabsch 2010): |m2 . (S x S0)| / (|S| |S0|) = zeta * sin(2theta).
            // Dividing by the scalar product S.S0 = |S||S0|cos(2theta) would add a spurious
            // 1/cos(2theta) to the absolute scale.
            float S_len = sqrtf(dot3(Sx, Sy, Sz, Sx, Sy, Sz));
            float S0_len = sqrtf(dot3(C.S0.x, C.S0.y, C.S0.z, C.S0.x, C.S0.y, C.S0.z));
            float lorentz = fabsf(dot3(C.m2.x, C.m2.y, C.m2.z, cx, cy, cz)) / (S_len * S0_len);

            // Partiality calculation (Kabsch formulation)
            // c1 = sqrt(2) * sigma / zeta, where sigma = mosaicity
            float c1 = zeta_abs / (sqrtf(2.0f) * mos_eff_rad);
            float half_wedge = C.wedge_angle_rad / 2.0f;
            float partiality = RockingSlicePartiality(phi, half_wedge, c1);


            // The partial sits at the centre of the frame's slice of its rocking curve, walked along
            // the Debye ring from the exact condition: S turned about the beam (see the CPU predictor).
            float kx = C.S0.x, ky = C.S0.y, kz = C.S0.z;
            normalize3(kx, ky, kz);
            const float kS = dot3(kx, ky, kz, Sx, Sy, Sz);
            const float qx = Sx - kS * kx, qy = Sy - kS * ky, qz = Sz - kS * kz;   // S_perp
            float tx, ty, tz;                                                   // S_turn = k x S
            cross3(kx, ky, kz, Sx, Sy, Sz, tx, ty, tz);
            float wx, wy, wz;                                                   // m2 x p
            cross3(C.m2.x, C.m2.y, C.m2.z, px, py, pz, wx, wy, wz);
            const float c_slice = RockingSliceCentroid_rad(phi, half_wedge, c1, partiality);
            const float psi = c_slice * dot3(wx, wy, wz, tx, ty, tz) / dot3(qx, qy, qz, qx, qy, qz);
            const float cos_psi = cosf(psi), sin_psi = sinf(psi);
            Sx = kS * kx + qx * cos_psi + tx * sin_psi;
            Sy = kS * ky + qy * cos_psi + ty * sin_psi;
            Sz = kS * kz + qz * cos_psi + tz * sin_psi;

            // Use S (rotated) for projection
            float Srx = C.rot[0] * Sx + C.rot[1] * Sy + C.rot[2] * Sz;
            float Sry = C.rot[3] * Sx + C.rot[4] * Sy + C.rot[5] * Sz;
            float Srz = C.rot[6] * Sx + C.rot[7] * Sy + C.rot[8] * Sz;

            if (Srz <= 0.0f) continue;

            float coeff = C.coeff_const / Srz;
            float x = C.beam_x + Srx * coeff;
            float y = C.beam_y + Sry * coeff;

            if (x < C.det_x0 || x >= C.det_x1 || y < C.det_y0 || y >= C.det_y1)
                continue;

            float dist_ewald = fabsf(sqrtf(Sx * Sx + Sy * Sy + Sz * Sz) - C.one_over_wavelength);

            // Sensor quantum efficiency and the air flight path at this reflection's angle of
            // incidence on the detector. Sr is the diffracted direction in the detector's own
            // frame, so Srz over its length is the cosine to the detector NORMAL - which carries
            // detector tilt for free. Mirrors sensor_absorption::SensorQE::Factor and
            // ::FlightPathAttenuation::Factor on the CPU side; the air term runs the other way,
            // because an oblique reflection crossed more air.
            float qe_corr = 1.0f;
            float flight_corr = 1.0f;
            {
                float cos_alpha = Srz / sqrtf(Sx * Sx + Sy * Sy + Sz * Sz);
                if (cos_alpha > 1e-3f) {
                    if (C.qe_a0 > 0.0f) {
                        float qe = 1.0f - expf(-C.qe_a0 / cos_alpha);
                        if (qe > 0.0f) qe_corr = C.qe_qe0 / qe;
                    }
                    if (C.flight_dL > 0.0f)
                        flight_corr = expf(C.flight_dL * (1.0f / cos_alpha - 1.0f));
                }
            }

            out[count].h = h;
            out[count].k = k;
            out[count].l = l;
            out[count].delta_phi_deg = phi * 180.0 / PI;
            out[count].predicted_x = x;
            out[count].predicted_y = y;
            out[count].observed_x = NAN;
            out[count].observed_y = NAN;
            out[count].d = 1.0f / sqrtf(p0_sq);
            out[count].dist_ewald = dist_ewald;
            out[count].prescaling_corr = lorentz;
            out[count].qe_corr = qe_corr;
            out[count].flight_corr = flight_corr;
            out[count].partiality = partiality;
            out[count].zeta = zeta_abs;
            out[count].image_scale_corr = lorentz * qe_corr * flight_corr / partiality;
            count++;
        }
        return count;
    }

    // The l of each (h, k) column that can pass the rocking-curve test (RotationColumnRange, as the CPU
    // predictor walks them), one thread a column. In double, which a consumer card runs at a small
    // fraction of the float rate and shares between the warps of a multiprocessor: so few operations a
    // column, and small blocks to spread the columns over the multiprocessors. phi_limit <= 0 walks every
    // l. A column with no l to walk gets l_lo > l_hi; an empty skip interval is moved out of the column.
    __global__ void column_range_kernel(const KernelConstsRot *__restrict__ kc, int max_h, int max_k, int max_l,
                                        RotationColumn *__restrict__ columns) {
        const int column = blockIdx.x * blockDim.x + threadIdx.x;
        if (column >= (2 * max_h + 1) * (2 * max_k + 1)) return;
        const KernelConstsRot &C = *kc;
        RotationColumn col{-max_l, max_l, max_l + 1, max_l + 1};
        if (C.phi_limit > 0.0) {
            const double h = column / (2 * max_k + 1) - max_h;
            const double k = column % (2 * max_k + 1) - max_k;
            const double b = C.b0 + C.b_h * h + C.b_k * k;
            const double c = (C.c_hh * h + C.c_hk * k + C.c_h) * h + (C.c_kk * k + C.c_k) * k;
            if (!RotationColumnRange(C.C_C, C.inv_2C_C, b, c, C.f_max, max_l, col))
                col = {max_l + 1, max_l, max_l + 1, max_l + 1};
            if (col.skip_lo > col.skip_hi)
                col.skip_lo = col.skip_hi = max_l + 1;
        }
        columns[column] = col;
    }

    // One (h, k) column per warp (a row of the block), its l spread over the warp's threads. The order
    // the atomics place the reflections in does not matter, as OrderAndSelect sorts them.
    __global__ void bragg_rot_kernel_columns(const KernelConstsRot *__restrict__ kc,
                                             const RotationColumn *__restrict__ columns,
                                             int max_h, int max_k,
                                             int max_reflections,
                                             Reflection *__restrict__ out,
                                             int *__restrict__ counter) {
        const int column = blockIdx.x * blockDim.y + threadIdx.y;
        if (column >= (2 * max_h + 1) * (2 * max_k + 1)) return;
        const int h = column / (2 * max_k + 1) - max_h;
        const int k = column % (2 * max_k + 1) - max_k;
        const RotationColumn col = columns[column];

        // The column's l as one list: [l_lo, skip_lo) and then (skip_hi, l_hi].
        const int first_end = min(col.l_hi, col.skip_lo - 1);
        const int n_first = max(0, first_end - col.l_lo + 1);
        const int second_start = max(col.l_lo, col.skip_hi + 1);
        const int n = n_first + max(0, col.l_hi - second_start + 1);

        for (int j = threadIdx.x; j < n; j += blockDim.x) {
            const int l = j < n_first ? col.l_lo + j : second_start + (j - n_first);
            Reflection r[2];
            const int found = compute_reflections_rot(*kc, h, k, l, r);
            for (int i = 0; i < found; ++i) {
                // Do NOT clamp the counter back down on overflow: it then saturates at the capacity and the
                // host cannot tell a full buffer from an overflowing one. Let it count the true total.
                const int pos = atomicAdd(counter, 1);
                if (pos < max_reflections)
                    out[pos] = r[i];
            }
        }
    }

    // A float as an unsigned integer that sorts in the same order. -0 is first made +0, which the
    // comparators of OrderOutput and TruncateToOutput treat as equal to it.
    __device__ inline uint32_t FloatOrderKey(float v) {
        const uint32_t u = __float_as_uint(v + 0.0f);
        return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
    }

    struct HklPacking {
        int max_h, max_k, max_l;
        int shift_h, shift_k;   // l takes the lowest bits
    };

    __global__ void phi_key_kernel(const Reflection *__restrict__ r, int n,
                                   uint32_t *__restrict__ key, int *__restrict__ index) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        key[i] = FloatOrderKey(r[i].delta_phi_deg);
        index[i] = i;
    }

    __global__ void hkl_key_kernel(const Reflection *__restrict__ r, const int *__restrict__ index, int n,
                                   HklPacking p, uint64_t *__restrict__ key) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        const Reflection &x = r[index[i]];
        key[i] = (static_cast<uint64_t>(x.h + p.max_h) << p.shift_h)
                 | (static_cast<uint64_t>(x.k + p.max_k) << p.shift_k)
                 | static_cast<uint64_t>(x.l + p.max_l);
    }

    // Largest partiality first, then smallest excitation error (see TruncateToOutput).
    __global__ void best_recorded_key_kernel(const Reflection *__restrict__ r, const int *__restrict__ index,
                                             int n, uint64_t *__restrict__ key) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        const Reflection &x = r[index[i]];
        key[i] = (static_cast<uint64_t>(~FloatOrderKey(x.partiality)) << 32) | FloatOrderKey(x.dist_ewald);
    }

    __global__ void gather_kernel(const Reflection *__restrict__ r, const int *__restrict__ index, int n,
                                  Reflection *__restrict__ out) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        out[i] = r[index[i]];
    }

    // Bits that hold 0 .. 2 * max_index.
    int IndexBits(int max_index) {
        int bits = 1;
        while ((1ll << bits) < 2ll * max_index + 1)
            bits++;
        return bits;
    }

    inline KernelConstsRot BuildKernelConstsRot(const DiffractionExperiment &experiment,
                                                const CrystalLattice &lattice,
                                                const BraggPredictionSettings &settings) {
         KernelConstsRot kc{};
         auto geom = experiment.GetDiffractionGeometry();

         const PlaneBox bounds = experiment.GetModules().PlaneBounds(experiment.GetXPixelsNum(), experiment.GetYPixelsNum());
         kc.det_x0 = static_cast<float>(bounds.x0);
         kc.det_y0 = static_cast<float>(bounds.y0);
         kc.det_x1 = static_cast<float>(bounds.x1);
         kc.det_y1 = static_cast<float>(bounds.y1);
         kc.beam_x = geom.GetBeamX_pxl();
         kc.beam_y = geom.GetBeamY_pxl();
         kc.coeff_const = geom.GetDetectorDistance_mm() / geom.GetPixelSize_mm();

         float one_over_dmax = 1.0f / settings.high_res_A;
         kc.one_over_dmax_sq = one_over_dmax * one_over_dmax;
         kc.one_over_wavelength = 1.0f / geom.GetWavelength_A();

         // Store mosaicity and wedge in radians for partiality calculation
         kc.mos_angle_rad = settings.mosaicity_deg * static_cast<float>(PI) / 180.0f;
         kc.wedge_angle_rad = settings.wedge_deg * static_cast<float>(PI) / 180.0f;
         kc.min_zeta = settings.min_zeta;
         kc.mosaicity_multiplier = settings.mosaicity_multiplier;
         kc.bandwidth_sigma = settings.bandwidth_sigma;
         kc.half_wavelength_A = geom.GetWavelength_A() / 2.0f;

         const auto &det = experiment.GetDetectorSetup();
         const auto sensor_qe = sensor_absorption::SensorQE::Build(
             det.GetSensorMaterial(), det.GetSensorThickness_um(), geom.GetWavelength_A());
         kc.qe_a0 = sensor_qe.active ? sensor_qe.a0 : 0.0f;
         kc.qe_qe0 = sensor_qe.qe0;
         kc.flight_dL = sensor_absorption::FlightPathAttenuation::Build(
                 experiment.GetBraggIntegrationSettings().GetFlightPath(),
                 geom.GetDetectorDistance_mm(), geom.GetWavelength_A()).d_over_L;

         kc.Astar = lattice.Astar();
         kc.Bstar = lattice.Bstar();
         kc.Cstar = lattice.Cstar();
         kc.S0 = geom.GetScatteringVector();

         auto rotT = geom.GetDetectorMatrix().transpose().arr();
         for (int i = 0; i < 9; ++i) kc.rot[i] = rotT[i];

         kc.centering = settings.centering;
         return kc;
     }

    inline void BuildGoniometerBasis(const DiffractionExperiment &experiment, KernelConstsRot &kc) {
        const auto gon_opt = experiment.GetGoniometer();
        if (!gon_opt.has_value())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "BraggPredictionRotationGPU requires a goniometer axis");
        const GoniometerAxis &gon = *gon_opt;

        // m2 = normalize(axis)
        float m2x = gon.GetAxis().x;
        float m2y = gon.GetAxis().y;
        float m2z = gon.GetAxis().z;
        normalize3(m2x, m2y, m2z);

        // m1 = normalize(m2 x S0)
        float m1x, m1y, m1z;
        cross3(m2x, m2y, m2z, kc.S0.x, kc.S0.y, kc.S0.z, m1x, m1y, m1z);
        normalize3(m1x, m1y, m1z);

        // m3 = normalize(m1 x m2)
        float m3x, m3y, m3z;
        cross3(m1x, m1y, m1z, m2x, m2y, m2z, m3x, m3y, m3z);
        normalize3(m3x, m3y, m3z);

        kc.m1 = Coord(m1x, m1y, m1z);
        kc.m2 = Coord(m2x, m2y, m2z);
        kc.m3 = Coord(m3x, m3y, m3z);
        kc.m2_S0 = dot3(m2x, m2y, m2z, kc.S0.x, kc.S0.y, kc.S0.z);
        kc.m3_S0 = dot3(m3x, m3y, m3z, kc.S0.x, kc.S0.y, kc.S0.z);
    }

    // The bound on the l of a column, as the CPU predictor sets it up (BraggPredictionRot.cpp), with the
    // column's terms as polynomials in h and k: b = 2 (AB.C* + S0.C*), c = AB.AB + 2 S0.AB, AB = h A* + k B*.
    inline void BuildColumnBound(const BraggPredictionSettings &settings, KernelConstsRot &kc) {
        const float one_over_dmax = 1.0f / settings.high_res_A;
        kc.phi_limit = RotationPhiLimit(settings.min_zeta, kc.mos_angle_rad, kc.bandwidth_sigma, kc.half_wavelength_A,
                                        one_over_dmax, kc.wedge_angle_rad / 2.0f, settings.mosaicity_multiplier);
        auto dot = [](const Coord &a, const Coord &b) {
            return static_cast<double>(a.x) * b.x + static_cast<double>(a.y) * b.y + static_cast<double>(a.z) * b.z;
        };
        const double S0_perp = std::sqrt(std::max(0.0, dot(kc.S0, kc.S0) - static_cast<double>(kc.m2_S0) * kc.m2_S0));
        kc.f_max = 1.01 * 2.0 * S0_perp * one_over_dmax * kc.phi_limit;
        kc.C_C = dot(kc.Cstar, kc.Cstar);
        kc.inv_2C_C = 1.0 / (2.0 * kc.C_C);
        kc.b0 = 2.0 * dot(kc.S0, kc.Cstar);
        kc.b_h = 2.0 * dot(kc.Astar, kc.Cstar);
        kc.b_k = 2.0 * dot(kc.Bstar, kc.Cstar);
        kc.c_hh = dot(kc.Astar, kc.Astar);
        kc.c_hk = 2.0 * dot(kc.Astar, kc.Bstar);
        kc.c_kk = dot(kc.Bstar, kc.Bstar);
        kc.c_h = 2.0 * dot(kc.S0, kc.Astar);
        kc.c_k = 2.0 * dot(kc.S0, kc.Bstar);
    }
} // namespace

BraggPredictionRotGPU::BraggPredictionRotGPU(int max_reflections)
    : BraggPrediction(0), dK(1), d_count(1), h_count(1) {
    GrowCapacity(max_reflections);
}

// Only the device side follows the prediction count: the host gets the kept rows alone, so
// `reflections` is sized in OrderAndSelect. Grown with slack, so a frame predicting a few more than
// the last record - the supercell probe's frames above all - does not reallocate every buffer and
// predict again.
void BraggPredictionRotGPU::GrowCapacity(int count) {
    const int capacity = count + count / 2;
    max_reflections = capacity;
    // Queued work may still use the buffers about to be freed (see CudaDevicePtr).
    cudaStreamSynchronize(stream);
    d_out = CudaDevicePtr<Reflection>(capacity);
    for (int i = 0; i < 2; i++) {
        d_key32[i] = CudaDevicePtr<uint32_t>(capacity);
        d_key64[i] = CudaDevicePtr<uint64_t>(capacity);
        d_index[i] = CudaDevicePtr<int>(capacity);
    }
    size_t bytes32 = 0, bytes64 = 0;
    cub::DoubleBuffer<uint32_t> k32(d_key32[0], d_key32[1]);
    cub::DoubleBuffer<uint64_t> k64(d_key64[0], d_key64[1]);
    cub::DoubleBuffer<int> idx(d_index[0], d_index[1]);
    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes32, k32, idx, capacity, 0, 32, stream));
    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes64, k64, idx, capacity, 0, 64, stream));
    sort_temp_bytes = std::max(bytes32, bytes64);
    d_sort_temp = CudaDevicePtr<uint8_t>(sort_temp_bytes);
}

int BraggPredictionRotGPU::OrderAndSelect(int count, const BraggPredictionSettings &settings) {
    const int kept = std::min(count, output_limit);
    if (kept > kept_capacity) {
        cudaStreamSynchronize(stream);
        kept_capacity = kept;
        d_kept = CudaDevicePtr<Reflection>(kept);
        h_out = CudaHostPtr<Reflection>(kept);
    }
    if (reflections.size() < static_cast<size_t>(kept))
        reflections.resize(kept);

    const int threads = 256;
    const int blocks = (count + threads - 1) / threads;
    cub::DoubleBuffer<uint32_t> k32(d_key32[0], d_key32[1]);
    cub::DoubleBuffer<uint64_t> k64(d_key64[0], d_key64[1]);
    cub::DoubleBuffer<int> idx(d_index[0], d_index[1]);
    size_t temp_bytes = sort_temp_bytes;

    // OrderOutput's order - h, k, l, then delta_phi - as two stable radix sorts, the minor key first.
    phi_key_kernel<<<blocks, threads, 0, stream>>>(d_out, count, k32.Current(), idx.Current());
    cuda_err(cudaGetLastError());
    cuda_err(cub::DeviceRadixSort::SortPairs(d_sort_temp, temp_bytes, k32, idx, count, 0, 32, stream));

    HklPacking pack{settings.max_h, settings.max_k, settings.max_l};
    const int bits_l = IndexBits(settings.max_l);
    const int bits_k = IndexBits(settings.max_k);
    pack.shift_k = bits_l;
    pack.shift_h = bits_l + bits_k;
    const int hkl_bits = pack.shift_h + IndexBits(settings.max_h);
    hkl_key_kernel<<<blocks, threads, 0, stream>>>(d_out, idx.Current(), count, pack, k64.Current());
    cuda_err(cudaGetLastError());
    cuda_err(cub::DeviceRadixSort::SortPairs(d_sort_temp, temp_bytes, k64, idx, count, 0, hkl_bits, stream));

    // Over the cap, TruncateToOutput's order on top: a stable sort leaves hkl breaking its ties.
    if (count > output_limit) {
        best_recorded_key_kernel<<<blocks, threads, 0, stream>>>(d_out, idx.Current(), count, k64.Current());
        cuda_err(cudaGetLastError());
        cuda_err(cub::DeviceRadixSort::SortPairs(d_sort_temp, temp_bytes, k64, idx, count, 0, 64, stream));
    }

    gather_kernel<<<(kept + threads - 1) / threads, threads, 0, stream>>>(d_out, idx.Current(), kept, d_kept);
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(h_out, d_kept, sizeof(Reflection) * kept, cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
    std::copy(h_out.get(), h_out.get() + kept, reflections.begin());
    return kept;
}

int BraggPredictionRotGPU::Calc(const DiffractionExperiment &experiment,
                                const CrystalLattice &lattice,
                                const BraggPredictionSettings &settings) {
    KernelConstsRot hK = BuildKernelConstsRot(experiment, lattice, settings);
    BuildGoniometerBasis(experiment, hK);
    BuildColumnBound(settings, hK);
    if (enumerate_full_box)
        hK.phi_limit = -1.0;

    cudaMemcpyAsync(dK, &hK, sizeof(KernelConstsRot), cudaMemcpyHostToDevice, stream);
    cudaMemsetAsync(d_count, 0, sizeof(int), stream);

    // Every (h, k) column, inclusive on both ends as the CPU loops (-max_i .. +max_i): the l of each to
    // walk first, then a warp's worth of l per column, eight columns per block.
    const int columns = (2 * settings.max_h + 1) * (2 * settings.max_k + 1);
    if (columns > column_capacity) {
        cudaStreamSynchronize(stream);
        column_capacity = columns;
        d_columns = CudaDevicePtr<RotationColumn>(columns);
    }
    column_range_kernel<<<(columns + 31) / 32, 32, 0, stream>>>(dK, settings.max_h, settings.max_k, settings.max_l,
                                                                 d_columns);
    cuda_err(cudaGetLastError());
    dim3 block(32, 8);
    dim3 grid((columns + block.y - 1) / block.y);

    bragg_rot_kernel_columns<<<grid, block, 0, stream>>>(dK, d_columns, settings.max_h, settings.max_k, max_reflections, d_out, d_count);
    cuda_err(cudaGetLastError());

    cudaMemcpyAsync(h_count, d_count, sizeof(int), cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);

    int count = *h_count.get();
    if (count > max_reflections) {
        // The buffer holds an arbitrary subset of what was predicted (whichever slots the atomics
        // reached first), so it cannot be used. Grow to fit and predict again; the buffer stays grown,
        // so a run pays for this a handful of times at most.
        GrowCapacity(count);
        cudaMemsetAsync(d_count, 0, sizeof(int), stream);
        bragg_rot_kernel_columns<<<grid, block, 0, stream>>>(dK, d_columns, settings.max_h, settings.max_k, max_reflections, d_out, d_count);
        cuda_err(cudaGetLastError());
        cudaMemcpyAsync(h_count, d_count, sizeof(int), cudaMemcpyDeviceToHost, stream);
        cudaStreamSynchronize(stream);
        count = std::min(*h_count.get(), max_reflections);
    }
    if (count == 0) return 0;
    return OrderAndSelect(count, settings);
}

#endif