Jungfraujoch/image_analysis/bragg_prediction/BraggPredictionRotGPU.h

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


#pragma once

#include <vector>
#include "BraggPrediction.h"
#include "../indexing/CUDAMemHelpers.h"

struct KernelConstsRot {
    float det_width_pxl;
    float det_height_pxl;
    float beam_x;
    float beam_y;
    float coeff_const;
    float one_over_wavelength;
    float one_over_dmax_sq;
    float mos_angle_rad;
    float wedge_angle_rad;
    float min_zeta;
    float mosaicity_multiplier;
    Coord Astar, Bstar, Cstar, S0;
    Coord m1, m2, m3;
    float m2_S0;
    float m3_S0;
    float rot[9];
    char centering;
};

class BraggPredictionRotGPU : public BraggPrediction {
    CudaRegisteredVector<Reflection> reg_out; // requires stable storage
    CudaDevicePtr<Reflection> d_out;

    // Dedicated stream
    CudaStream stream;

    // Device allocations via helpers
    CudaDevicePtr<KernelConstsRot> dK;
    CudaDevicePtr<int> d_count;

    // Host pinned buffer for async download (optional but faster)
    CudaHostPtr<int> h_count;

public:
    explicit BraggPredictionRotGPU(int max_reflections = kPredictionCapacity);

    int Calc(const DiffractionExperiment &experiment,
             const CrystalLattice &lattice,
             const BraggPredictionSettings &settings) override;
};