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

#pragma once

#include <array>
#include <chrono>
#include <cstdint>
#include <complex>
#include <memory>
#include <optional>
#include <vector>

#include "gemmi/grid.hpp"
#include "gemmi/model.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

#include "../../common/Reflection.h"   // MergedReflection
#include "RefinementRestraints.h"     // RestraintSet
#include "RigidBodyRefine.h"           // SymmetryComposition

class Logger;

// The X-ray side of restrained refinement (Refinement.h): the observed amplitudes, the model's structure
// factors and the maximum-likelihood target that compares them, with its gradient with respect to every
// atom's position and B.

// The reflections a refinement fits and scores: the observed amplitudes, in the group's reciprocal ASU, that
// the refinement grid composes (SymmetryComposition) - every observed one in practice.
struct RefinementReflections {
    std::vector<gemmi::Miller> hkl;
    std::vector<float> fo, sigma;
    std::vector<char> free, centric;
    std::vector<float> epsilon;
    std::vector<double> inv_d2;
    size_t Size() const { return hkl.size(); }
};

// The French-Wilson amplitudes of `merged` (NaN dropped) into the ASU of `sg`, those the grid of
// RigidBodyZoneGrid(cell, sg, d_min) composes.
RefinementReflections PrepareRefinementReflections(const std::vector<MergedReflection> &merged,
                                                   const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                   double d_min);

// What turns the atoms' structure factors F into the model the target compares with the data, per reflection:
//   F_model = k_total (F + fmask),
// with k_total the overall scale and anisotropic B and fmask the bulk solvent already scaled by k_sol
// exp(-b_sol s^2/4); and the two parameters of the likelihood, D and the variance. Fixed for one macro-cycle.
struct XrayRowModel {
    std::vector<float> k_total;
    std::vector<std::complex<float>> fmask;
    std::vector<float> d, var;
    // D per likelihood bin, where EstimateLikelihoodParameters() found it - the next fit's start (empty: none)
    std::vector<double> bin_d;
};

// The flat bulk-solvent model's mask (SolventMasker with the Refmac radii, PutMaskOnGrid) and its structure factors
// at the reflections, on a grid of its own rather than the atoms' density grid: the mask's edge is set by the
// probe and shrink radii (1.0 and 0.8 A), which a grid at d_min / 3 cannot resolve at low resolution - at 3 A
// its 1 A step leaves no grid point within the shrink radius and the shrink is skipped. The step is therefore
// d_min / 3 or MASK_SPACING_MAX, whichever is finer - where the mask grid stays within MASK_POINTS_MAX points; a
// cell too large for that gets the spacing that fills it. The R-factors of low-resolution models settle at a
// step of 0.4 A (within 0.001 of 0.3 A) and move by up to 0.005 between 0.4 and 0.8 A.
// The flat solvent model of Fokine & Urzhumtsev (2002) Acta Cryst. D58, 1387-1392
// The probe-and-shrink mask of Jiang & Brunger (1994) J. Mol. Biol. 243, 100-115
class BulkSolventMask {
public:
    static constexpr double MASK_SPACING_MAX = 0.4;   // A
    static constexpr double MASK_POINTS_MAX = 1 << 26;
    BulkSolventMask(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, size_t nthreads);
    // The mask's grid (sizes and spacing set, values not allocated) - also the GPU engine's (RefinementGPU.cpp).
    static gemmi::Grid<float> MaskGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min);
    // Fmask (unscaled) of `model` at `hkl`.
    void Fmask(const gemmi::Model &model, const std::vector<gemmi::Miller> &hkl,
               std::vector<std::complex<float>> &f) const;
    double Spacing() const { return spacing_; }

private:
    gemmi::Grid<float> grid_;
    mutable std::vector<size_t> orbit_leaders_;   // made at the first Fmask()
    size_t nthreads_;
    double spacing_;
};

// The ML parameters per reflection, in resolution bins of the free set interpolated between the bins' centres:
// beta fitted by maximum likelihood to the free reflections as the whole spread of Fo about D Fc, D then fitted to
// the working reflections with the variance epsilon beta + 2 sigma^2 (acentric; + sigma^2 centric), which is also
// the variance of a reflection in the target. `fmodel` is k_total (F + fmask), without D.
// Following Lunin & Skovoroda (1995) Acta Cryst. A51, 880-887, and Murshudov et al. (2011) Acta Cryst. D67, 355-367
void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel,
                                  XrayRowModel &rows, size_t nthreads = 1);

// R-work and R-free of k_total (F + fmask) against the observations.
struct RFactors {
    double r_work = 1, r_free = 1;
    int n_work = 0, n_free = 0;
};
RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel);

// x, y, z and B of every atom of `model`, in model order.
std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model);

// What evaluates the targets for a set of atom positions and B: the CPU below, or a GPU (RefinementGPU.h). The
// atoms are the model's, all of them, in model order; the grid and the reflections are fixed when the engine is
// made, the row model each macro-cycle, the blur where SetBlur is called.
class RefinementEngine {
public:
    virtual ~RefinementEngine() = default;
    // The blur the atoms are gridded with (DensityCalculator's), and the matching unblur of their transform.
    // Every atom's B must stay above what the blur was chosen for (RefinementBlur) while it is in use.
    virtual void SetBlur(double blur) = 0;
    virtual void SetRowModel(const XrayRowModel &rows) = 0;
    // F of the atoms at the reflections.
    virtual void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) = 0;
    // The bulk-solvent mask's structure factors at the reflections (unscaled).
    virtual void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) = 0;
    // The sum of -ln L over the working reflections and, with `grad`, its derivatives with respect to every
    // atom's x, y, z (Cartesian, A) and B (A^2), in model order.
    virtual double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) = 0;
    virtual bool OnGPU() const = 0;

    // Restraint set `which` (0: the geometry and the repulsions, 1: the B similarity) and its sum of E at xyzb
    // (x, y, z, B per atom) with, given `grad`, dE/d(x, y, z, B) per atom. On the CPU unless the engine has its own.
    virtual void SetRestraints(int which, const RestraintSet &set) { restraints_[which] = set; }
    virtual double EvaluateRestraints(int which, const std::vector<std::array<double, 4>> &xyzb,
                                      std::vector<std::array<double, 4>> *grad) {
        if (grad != nullptr)
            grad->assign(xyzb.size(), {0, 0, 0, 0});
        return ::EvaluateRestraints(restraints_[which], xyzb, grad);
    }
    // Evaluate() and EvaluateRestraints(which) at the atoms of `model` as they are, in one go: returns the X-ray
    // target, and the restraints' sum in `restraints`.
    virtual double EvaluateWithRestraints(const gemmi::Model &model, int which, std::vector<std::array<double, 4>> *grad,
                                          std::vector<std::array<double, 4>> *restraint_grad, double &restraints) {
        const double xray = Evaluate(model, grad);
        restraints = EvaluateRestraints(which, ModelXyzB(model), restraint_grad);
        return xray;
    }

private:
    RestraintSet restraints_[2];
};

// The time a refinement spends in one stage, summed over every call (and over threads) until Log(): the scope of
// an instance is one call, and an NVTX range of the stage's name.
class RefinementStage {
public:
    explicit RefinementStage(const char *name);
    ~RefinementStage();
    RefinementStage(const RefinementStage &) = delete;
    RefinementStage &operator=(const RefinementStage &) = delete;
    // "TIMING refine.<stage> <seconds> <calls>" per stage at debug level, then the sums start again.
    static void Log(Logger &logger);

private:
    const char *name_;
    std::chrono::steady_clock::time_point start_;
    uint64_t nvtx_;
};

// The blur for a model whose smallest B may go as low as b_floor: DensityCalculator's
// set_refmac_compatible_blur() for that B, on the refinement grid of d_min.
double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor);

// The CPU engine: the density of one copy of the atoms on the grid of RigidBodyZoneGrid(cell, sg, d_min)
// (PutModelDensityOnGrid), its transform and the crystal's symmetry composed in reciprocal space
// (SymmetryComposition). The gradient by the FFT method: the derivative of the target with respect to each
// structure factor, taken back to real space by one inverse transform, is a map; each atom's derivatives are
// that map summed against the derivatives of its own density.
// Following Agarwal (1978) Acta Cryst. A34, 791-809
std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell &cell,
                                                          const gemmi::SpaceGroup &sg, double d_min,
                                                          const RefinementReflections &refl, size_t nthreads);

// The derivative map of the X-ray target on a grid (the FFT method's map), summed against each atom's density
// derivatives: dT/d(x, y, z, B) per atom. `map` is u fastest on the grid of RigidBodyZoneGrid(), the atoms
// gridded with `blur`. Exposed for the tests.
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
                          std::vector<std::array<double, 4>> &grad);