// 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 KernelConsts {
    float det_x0, det_y0, det_x1, det_y1;   // the plane the panels span (DetectorModules::PlaneBounds)
    float beam_x;
    float beam_y;
    float coeff_const;
    float one_over_wavelength;
    float one_over_dmax_sq;
    float ewald_cutoff;
    float bandwidth_sigma;   // relative Δλ/λ (sigma); 0 = monochromatic
    // Angle-dependent sensor efficiency, reduced to the two per-dataset numbers the kernel needs:
    // the optical thickness at normal incidence and the absorbed fraction there. qe_a0 = 0 leaves
    // every reflection untouched (see sensor_absorption::SensorQE).
    float qe_a0;
    float qe_qe0;
    float flight_dL;   // normal-incidence air flight path in attenuation lengths; 0 = off
    Coord Astar, Bstar, Cstar, S0;
    float rot[9];
    char centering;
};

class BraggPredictionGPU : public BraggPrediction {
    // The device's reflections land here and are copied into `reflections`: pinned, so the download is a
    // DMA, and from the pool, so the engine does not register and unregister a vector of its own.
    CudaHostPtr<Reflection> h_out;
    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 = kPredictionCapacity);
    // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
    ~BraggPredictionGPU() override { cudaStreamSynchronize(stream); }

protected:
    void GrowCapacity(int count) override;

public:

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