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

#pragma once

// The molecular-replacement search on a GPU (CUDA builds only; the header itself needs no CUDA): the
// rotation function and the translation function of MolecularReplacement.h, the same functions as the
// CPU computes - the same trilinear interpolation of the same molecular transform, the same bins, the
// same coefficients on the same grid - in float on the device, with cuFFT for FFTW. The two agree to
// rounding, not bit for bit. Every reduction runs in a fixed order and nothing is summed with an
// atomic, so the GPU is deterministic on its own.

#include <memory>
#include <vector>

#include "MolecularReplacement.h"

// MapToFPhi (ModelFFT.h) with cuFFT on the calling thread's current device: the same transform, the same
// half-l layout and conventions, to float rounding. For the molecular transform of a large model, whose box is
// hundreds of millions of points. MapToFPhiGPUFits: whether its device buffers take at most a quarter of the
// card's total memory - decided on the total, so the same input always takes the same path.
bool MapToFPhiGPUFits(const gemmi::Grid<float> &map);
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map);

struct MolecularReplacementGPUState;
struct MolecularReplacementGPUGroup;

class MolecularReplacementGPU {
public:
    // Uploads the molecular transform, on the calling thread's current device. Throws JFJochException on
    // a CUDA failure, as every call below does.
    explicit MolecularReplacementGPU(const MolecularTransform &ft);
    ~MolecularReplacementGPU();

    std::vector<double> RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                         const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations);

    // The translation function of one group at each of `rotations`; with `known`, beside a known partial
    // structure (KnownStructureFactors).
    std::vector<TranslationPeak> TranslationFunction(const TranslationSetup &setup,
                                                     const std::vector<gemmi::Mat33> &rotations,
                                                     const std::vector<std::complex<double>> *known = nullptr);

private:
    std::unique_ptr<MolecularReplacementGPUState> state_;
};