// 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 <complex>
#include <functional>
#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; `packed`: with each point's
    // interpolation corners together, read from two places instead of four, for more memory (TransformBytes)
    // and to the same bits. Throws JFJochException on a CUDA failure, as every call below does.
    explicit MolecularReplacementGPU(const MolecularTransform &ft, bool packed = false);
    ~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 the transform is held in, plain or packed; while the packed one is made, both.
    static size_t TransformBytes(const MolecularTransform &ft, bool packed);
    // 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_;
};

// The search on one card: the transform uploaded on `device`.
struct MolecularReplacementGPUOnCard {
    int device = 0;
    MolecularReplacementGPU *gpu = nullptr;
};

// What a function's slice holds on its card while it runs: called on the slice's thread, with the card and the
// bytes its buffers take, and kept until the slice is done. Empty: nothing is held.
using MolecularReplacementGPUHold = std::function<std::shared_ptr<void>(int device, size_t bytes)>;

// The rotation and translation functions with the orientations split over `cards`: the k-th card scores the
// k-th of cards.size() contiguous slices of `rotations`, on a thread of its own beside the others, and the
// results are put back in the order of `rotations`. Every orientation is scored by itself - one block of the
// rotation function, one map of the translation function - and nothing is summed across orientations, so the
// result is that of one card, bit for bit, on any number of cards.
std::vector<double> SplitRotationFunction(const std::vector<MolecularReplacementGPUOnCard> &cards,
                                          const SearchReflections &refl, const gemmi::UnitCell &cell,
                                          const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations,
                                          const MolecularReplacementGPUHold &hold = {});
std::vector<TranslationPeak> SplitTranslationFunction(const std::vector<MolecularReplacementGPUOnCard> &cards,
                                                      const TranslationSetup &setup,
                                                      const std::vector<gemmi::Mat33> &rotations,
                                                      const std::vector<std::complex<double>> *known,
                                                      TranslationSetupGPU *prepared,
                                                      const MolecularReplacementGPUHold &hold = {});