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

// CropHalfL(MapToFPhi(map), reach) (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. MapToFPhiGPUBytes: the device memory it takes while it runs.
size_t MapToFPhiGPUBytes(const gemmi::Grid<float> &map, int reach);
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map, int reach);

struct MolecularReplacementGPUState;
struct MolecularReplacementGPUGroup;

struct TranslationSetupGPU;

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). `prepared`: Prepare() of this setup, for a setup used again and again;
    // made for this call otherwise.
    std::vector<TranslationPeak> TranslationFunction(const TranslationSetup &setup,
                                                     const std::vector<gemmi::Mat33> &rotations,
                                                     const std::vector<std::complex<double>> *known = nullptr,
                                                     TranslationSetupGPU *prepared = nullptr);
    // What the translation function needs of `setup`, worked out once on the host; with `known_part`, for the
    // function beside a known partial structure.
    std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> Prepare(const TranslationSetup &setup,
                                                                                bool known_part);

    // The device memory a call takes for as long as it runs, and gives back when it returns: RotationFunction()
    // of `rotations` against `refl`'s reflections in `sg`, and TranslationFunction() of `rotations` (with
    // `known`, beside a known partial structure).
    static size_t RotationBytes(const SearchReflections &refl, const gemmi::SpaceGroup &sg, size_t rotations);
    static size_t TranslationBytes(const TranslationSetup &setup, size_t rotations, bool known);

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