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

#pragma once

// Restrained refinement's targets on a GPU (CUDA builds only): the RefinementEngine of RefinementXray.h with
// everything per evaluation on the device - the atoms' density (ModelDensityGPU's gather), its transform, the
// symmetry composed in reciprocal space, the likelihood target, the derivative map by the inverse transform
// of the target's derivatives and every atom's gradient read off it, and the geometry restraints. The host
// passes positions in and gets one value and one gradient per atom back.

#include <memory>

#include "RefinementXray.h"

// What an engine holds on its card: its own part; the part the engines like it on the same card share (the
// reflections' tables and the bulk-solvent mask's buffers), which the first of them there makes; and a lane of the
// transforms' working buffers, which an engine has to itself or shares with others, one evaluation at a time.
struct RefinementGPUBytes {
    size_t engine = 0, card = 0, lane = 0;
};

// The engine for refining `model` against `refl` on `device`, or null where the GPU cannot take this case (an
// atom's box wider than the cell, too many symmetry operators) or its memory is over `budget` bytes.
// `b_max`: the largest B an atom may reach, which the gather's reserve is sized for. Throws JFJochException on a
// CUDA failure. `nthreads`: the host threads that work out the atoms' densities. `restraint_records`,
// `restraint_slots`: the most records (plane atoms included) and gradient slots its two restraint sets hold together,
// which the memory is worked out for (0: about thirty records and sixty slots per atom).
std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell,
                                                          const gemmi::SpaceGroup &sg, double d_min,
                                                          const RefinementReflections &refl,
                                                          const gemmi::Model &model, double b_max, size_t budget,
                                                          RefinementGPUBytes *bytes = nullptr, size_t nthreads = 1,
                                                          size_t restraint_records = 0, size_t restraint_slots = 0);
// An engine like `like` (one the function above made) on `device`: the same case, and on a card where an engine
// like it is, that card's shared part with it - so it holds RefinementGPUBytes::engine there, and ::lane beside where
// `own_lane` (else it shares the card's least used lane). Throws JFJochException on a CUDA failure.
std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(const RefinementEngine &like, int device, size_t nthreads,
                                                          bool own_lane);