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

#ifndef JFJOCH_BRAGGPREDICTIONGPU_H
#define JFJOCH_BRAGGPREDICTIONGPU_H

#include <vector>

#include "BraggPrediction.h"

#include "../indexing/CUDAMemHelpers.h"

struct KernelConsts {
    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 ewald_cutoff;
    Coord Astar, Bstar, Cstar, S0;
    float rot[9];
    char centering;
};

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

    // Dedicated stream
    CudaStream stream;

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

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

public:
    explicit BraggPredictionGPU(int max_reflections = 10000);

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


#endif //JFJOCH_BRAGGPREDICTIONGPU_H