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

#pragma once

// This include should be only included in sections of the code, where it is certain that CUDA is present
// so with JFJOCH_USE_CUDA preprocessor definition, given this file is included in the source only in this case

#include <vector>
#include <mutex>
#include <optional>
#include "CUDAMemHelpers.h"
#include "../../common/Coord.h"
#include "../../common/CrystalLattice.h"
#include "FFTIndexer.h"
#include "../common/IndexingSettings.h"
#include "FFTResult.h"

// The histograms are made and transformed FFT_BATCH_DIRECTIONS directions at a time. All at once, the
// histograms, their transforms and cuFFT's work area held 0.8 GB at the default 16384 directions and 2 GB
// at a doubled cell bound, for the length of the indexer's life; a batch holds a quarter of that. A
// direction's histogram and transform do not depend on which batch it is in - cuFFT's R2C gives the same
// bits at any batch count (checked at the histogram sizes the indexer uses) - so neither does the result.
// Kept large enough that a batch still fills the card.
constexpr int FFT_BATCH_DIRECTIONS = 4096;

class FFTIndexerGPU : public FFTIndexer {
    CudaDevicePtr<float> d_dir_x;
    CudaDevicePtr<float> d_dir_y;
    CudaDevicePtr<float> d_dir_z;

    CudaDevicePtr<float> d_spot_x;
    CudaDevicePtr<float> d_spot_y;
    CudaDevicePtr<float> d_spot_z;

    CudaHostPtr<float> spot_x;
    CudaHostPtr<float> spot_y;
    CudaHostPtr<float> spot_z;

    // The histograms and their transforms of one batch of directions, and the magnitudes of all of them.
    CudaDevicePtr<float> d_input_fft;
    CudaDevicePtr<cufftComplex> d_output_fft;
    CudaDevicePtr<float> d_mag_t;
    CudaDevicePtr<FFTResult> d_result_fft;

    CudaRegisteredVector<FFTResult> result_fft_reg;

    CudaStream stream;

    // After the stream: it waits for the stream before the plan goes back, and goes before the stream.
    CudaFFTPlan plan;
    CudaFFTPlan plan_last;   // the last batch, where FFT_BATCH_DIRECTIONS does not divide the directions

    void ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) override;
    void DirectionsChanged() override;
public:
    explicit FFTIndexerGPU(const IndexingSettings& settings);
    FFTIndexerGPU(const FFTIndexerGPU &i) = delete;
    const FFTIndexerGPU &operator=(const FFTIndexerGPU &i) = delete;
    ~FFTIndexerGPU() override = default;
};