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

#pragma once

#include <algorithm>
#include <array>
#include <cstddef>
#include <iterator>
#include <utility>
#include <vector>

#include "gemmi/dencalc.hpp"   // DensityCalculator
#include "gemmi/grid.hpp"
#include "gemmi/it92.hpp"      // IT92 x-ray form factors
#include "gemmi/model.hpp"

// gemmi's gridding of a model - the density of DensityCalculator::put_model_density_on_grid() and the
// bulk-solvent mask of SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid() - spread over
// `nthreads` threads of ParallelFor's pool. Every grid point receives the same operations in the same
// order as in gemmi's serial loops, so the grids are the same on any number of threads: the mask is gemmi's
// bit for bit, the density gemmi's but for the last bits (its points are placed in single precision, as the
// Apple GPU places them).

// The points gemmi's Grid::symmetrize() reduces the orbits of this grid's size and space group into
// (empty in P1). Both functions below take them, for a grid of the same size.
std::vector<size_t> OrbitLeaders(const gemmi::Grid<float> &grid, size_t nthreads);

// The IT92 coefficients of an atom of element `el` and charge `charge`: the ion's where IT92 tabulates that ion
// (Mg2+, Na1+, Cl1- ...), the neutral atom's otherwise. gemmi's own lookup, IT92::get(), takes the neutral atom
// whatever the charge unless its global ignore_charge is turned off; this one does not depend on it.
template <class Real>
const typename gemmi::IT92<Real>::Coef &IonFormFactor(gemmi::El el, signed char charge) {
    using T = gemmi::IT92<Real>;
    if (charge != 0)
        for (size_t i = 0; i < std::size(T::ion_list); i++)
            if (T::ion_list[i] == std::pair<gemmi::El, signed char>(el, charge))
                return T::data[99 + i];
    return T::get(el, 0);
}

// The B an anisotropic atom's density is cut at the radius of: its largest eigenvalue (of 8 pi^2 U plus the blur), so
// that the density is cut where it falls to the cutoff along its widest axis. gemmi takes the largest diagonal element,
// which a U with large off-diagonal elements leaves short of that axis.
template <typename Real>
Real AnisoRadiusB(const gemmi::SMat33<Real> &aniso_b) {
    const gemmi::SMat33<double> m{aniso_b.u11, aniso_b.u22, aniso_b.u33, aniso_b.u12, aniso_b.u13, aniso_b.u23};
    const std::array<double, 3> ev = m.calculate_eigenvalues();
    return static_cast<Real>(std::max({ev[0], ev[1], ev[2], double(aniso_b.u11), double(aniso_b.u22), double(aniso_b.u33)}));
}

// `dc` set up as for put_model_density_on_grid(): d_min, rate, blur, the grid's cell and group. On macOS the
// Apple GPU puts the density on the grid (ModelDensityMetal.h), to the same bits; `gpu` false keeps it on the
// CPU. `ion_charges`: each atom scatters as the ion its charge makes it (IonFormFactor), not as gemmi's
// put_model_density_on_grid() takes it, neutral.
void PutModelDensityOnGrid(gemmi::DensityCalculator<gemmi::IT92<float>, float> &dc, const gemmi::Model &model,
                           const std::vector<size_t> &orbit_leaders, size_t nthreads, bool gpu = true,
                           bool ion_charges = false);

// `grid` sized, with its cell and group set.
void PutMaskOnGrid(gemmi::Grid<float> &grid, const gemmi::Model &model, const std::vector<size_t> &orbit_leaders,
                   size_t nthreads);

// The step a bulk-solvent mask's grid is held below. The mask's shrink (0.8 A, the Refmac radii) reaches no grid point
// on a coarser grid and is skipped, and the macromolecule keeps the whole 1.0 A probe: some 5-10% of the cell's solvent
// lost, 0.015-0.03 of R-free on deposited models. Below the shrink radius the R-factors no longer depend on the step.
constexpr double MASK_SPACING_SHRINK = 0.75;   // A

// The grid of a bulk-solvent mask for structure factors to d_min (sizes, cell and group set, no values): at `spacing`,
// but never coarser than MASK_SPACING_SHRINK nor than the density's own grid at rate 1.5 (d_min / 3), which it then
// is - set_size_from_spacing()'s size for d_min / 3.
gemmi::Grid<float> SolventMaskGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                                   double spacing);

// The number of points of the grid a model's structure factors to d_min are computed on at DensityCalculator's
// `rate` (a spacing of d_min / (2 rate)) in this cell and space group: the grid's size, without making it.
size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate = 1.5);

// What computing a model's structure factors on the CPU holds at its peak, per point of that grid: the density
// and the bulk-solvent mask grids, and for each of them its transform's input copy, spectrum and coefficients,
// with the two made at once (ModelValidation's compute_model_factors).
constexpr size_t MODEL_FACTORS_BYTES_PER_GRID_POINT = 32;
// What a molecular-replacement placement's structure factors and rigid body hold on the CPU at their peak, per
// point of that grid: measured at under half of the above.
constexpr size_t MODEL_PLACEMENT_BYTES_PER_GRID_POINT = 16;