// 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 "RotationColumnRange.h"
#include "../indexing/CUDAMemHelpers.h"

struct KernelConstsRot {
    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 mos_angle_rad;
    float wedge_angle_rad;
    float min_zeta;
    float mosaicity_multiplier;
    float bandwidth_sigma;      // relative dlambda/lambda as a sigma; 0 = monochromatic
    float half_wavelength_A;
    // 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;
    Coord m1, m2, m3;
    float m2_S0;
    float m3_S0;
    float rot[9];
    char centering;
    // Which l of a column the kernel walks (RotationColumnRange): phi_limit <= 0 walks them all. The
    // column's b(h, k) = b0 + b_h h + b_k k and c(h, k) = c_hh h^2 + c_hk h k + c_kk k^2 + c_h h + c_k k.
    double phi_limit, f_max, C_C, inv_2C_C;
    double b0, b_h, b_k, c_hh, c_hk, c_kk, c_h, c_k;
};

class BraggPredictionRotGPU : public BraggPrediction {
    // The device's reflections, in the order the atomics placed them.
    CudaDevicePtr<Reflection> d_out;
    // The rows Calc returns, ordered and selected on the device, then downloaded into h_out: pinned, so
    // the download is a DMA, and from the pool, so the engine does not register and unregister a vector
    // of its own. Sized for what is returned (at most output_limit), not for what was predicted.
    CudaDevicePtr<Reflection> d_kept;
    CudaHostPtr<Reflection> h_out;
    int kept_capacity = 0;

    // Sort keys and the permutation they carry, double-buffered for cub's radix sort, and its scratch.
    CudaDevicePtr<uint32_t> d_key32[2];
    CudaDevicePtr<uint64_t> d_key64[2];
    CudaDevicePtr<int> d_index[2];
    CudaDevicePtr<uint8_t> d_sort_temp;
    size_t sort_temp_bytes = 0;

    // Which l of each (h, k) column the kernel walks (column_range_kernel).
    CudaDevicePtr<RotationColumn> d_columns;
    int column_capacity = 0;

    // 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;

    // Order the `count` device reflections exactly as OrderOutput and TruncateToOutput would, keep the
    // first min(count, output_limit) and download them into `reflections`. Returns the kept count.
    int OrderAndSelect(int count, const BraggPredictionSettings &settings);

public:
    explicit BraggPredictionRotGPU(int max_reflections = kPredictionCapacity);
    // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
    ~BraggPredictionRotGPU() override { cudaStreamSynchronize(stream); }

protected:
    void GrowCapacity(int count) override;
    // Walk every l of every column, as BraggPredictionRot's flag of the same name - for the test.
    bool enumerate_full_box = false;

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