// 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(); }

    // With a twin law (SetTwinLaw) the structure factors are needed at every observed reflection's twin mate too:
    // the mates not observed follow the observed rows as rows of their own, with no observation (mate_hkl), and
    // `mate` holds per row, observed or not, the row of its mate - the row itself where the mate is the reflection
    // or is not composed. A law whose square is in the group makes it an involution: mate[mate[m]] == m.
    // Empty without a law.
    gemmi::Op twin_law = gemmi::Op::identity();
    std::vector<gemmi::Miller> mate_hkl;
    std::vector<int> mate;
    bool Twinned() const { return !mate.empty(); }
    // The observed rows, then the mates': the rows the structure factors are computed at.
    size_t Rows() const { return hkl.size() + mate_hkl.size(); }
    std::vector<gemmi::Miller> RowHkl() const;
};

// 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);

// The twin laws a refinement in `sg` considers: the rotations of the lattice's metric symmetry that `sg` lacks
// (Le Page two-folds within 3 degrees of obliquity, as AssignRfreeFlags groups the free set), one per coset of
// the group, each with its square in the group - a law that exchanges two domains. Miller-index operators in the
// data's basis.
// Following Lebedev, Vagin & Murshudov (2006) Acta Cryst. D62, 83-95
std::vector<gemmi::Op> RefinementTwinLaws(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg);

// `refl` with twin law `law` (RefinementReflections::mate). Returns how many twin pairs of observed reflections
// have one member in the working set and the other in the free set: with any, the free set is not independent
// of the working one under the law.
int SetTwinLaw(RefinementReflections &refl, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
               const gemmi::Op &law);

// 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.
// With a twin law, the model compared with an observation at h is the twinned one (TwinnedModel):
//   |F_model(h)|^2 = k_total(h)^2 ((1-a) |F(h) + fmask(h)|^2 + a |F(Th) + fmask(Th)|^2),
// a the twin fraction; fmask is then per row of refl.Rows(), the other per observed row.
struct XrayRowModel {
    std::vector<float> k_total;
    std::vector<std::complex<float>> fmask;
    std::vector<float> d, var;
    float twin_fraction = 0;
    // 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, but never one coarser than MASK_SPACING_SHRINK, where the
// shrink still acts, nor than the atoms' density grid (d_min / 3), which it then is (SolventMaskGrid, ModelGrid.h).
// 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; past the shrink radius there is no shrink at all, which costs far more.
// 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);
    // The step that grid is asked for (before good_grid_size() rounds it finer).
    static double MaskSpacing(const gemmi::UnitCell &cell, 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);

// F_model per observed row from F at every row: k_total (F + fmask), or with a twin law the twinned model - its
// amplitude, with the phase of the row's own k_total (F + fmask).
// Following Yeates (1997) Methods Enzymol. 276, 344-358 (I_obs(h) = (1-a) I(h) + a I(Th))
std::vector<std::complex<float>> TwinnedModel(const RefinementReflections &refl, const XrayRowModel &rows,
                                              const std::vector<std::complex<float>> &f);

// The twin fraction a in [0, 1/2] the working observations ask of the model: Fo^2 regressed, with an intercept, on
// the model's intensities at h and at Th, all normalised per resolution shell of TWIN_FRACTION_SHELL reflections,
// a = b(Th) / (b(h) + b(Th)). The expectation of Fo^2 given the model is linear in the two, D^2 |F_model|^2 + Sigma,
// so an unrelated reflection takes a coefficient of 0 however poor the model is - where a least-squares fit of the
// twinned amplitudes took a fraction of 0.16 for unrelated partners of one model at R-free 0.31, the average of two
// predictions being the better prediction. Reflections that are their own mate say nothing of a and are left out. `f` is F at every row, `rows`
// gives k_total and fmask.
// Following Yeates (1997) Methods Enzymol. 276, 344-358 (the twinned intensity), refined as Murshudov et al. (2011)
// Acta Cryst. D67, 355-367 refine the fractions, by least squares
constexpr size_t TWIN_FRACTION_SHELL = 500;
double FitTwinFraction(const RefinementReflections &refl, const XrayRowModel &rows,
                       const std::vector<std::complex<float>> &f);

// 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);

// An atom's displacement where it is anisotropic: U = dev + B / (8 pi^2) I, B the atom's b_iso (its equivalent
// isotropic B, as dev is traceless) and dev the anisotropy, U11 U22 U33 U12 U13 U23 (Cartesian, A^2). So everything
// that refines, restrains and bounds B works on an anisotropic atom as on an isotropic one, its shape held.
struct AtomAniso {
    bool on = false;
    std::array<float, 6> dev{};
};
// U of an anisotropic atom of b_iso `b`, its smallest eigenvalue at least ANISO_U_MIN (a U that is not positive
// definite has no density).
constexpr double ANISO_U_MIN = 0.5 / (8 * 3.14159265358979323846 * 3.14159265358979323846);   // B 0.5 A^2
gemmi::SMat33<float> AnisoU(const AtomAniso &a, double b);
// The model with every anisotropic atom of `aniso` given its U (AnisoU), the others as they are.
void ApplyAniso(gemmi::Model &model, const std::vector<AtomAniso> &aniso);

// 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 rows (refl.Rows(): the observed reflections, then the twin mates not observed).
    virtual void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) = 0;
    // The bulk-solvent mask's structure factors at the rows (unscaled).
    virtual void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) = 0;
    // The sum of -ln L over the working reflections - against the twinned model where refl has a twin law - and,
    // with `grad`, its derivatives with respect to every atom's x, y, z (Cartesian, A) and B (A^2), in model order.
    // For an anisotropic atom (SetAniso) the B derivative is the one along dev + B / (8 pi^2) I, its shape held.
    virtual double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) = 0;
    virtual bool OnGPU() const = 0;

    // Which atoms are anisotropic and their shapes, in model order (AtomAniso), read at every Fcalc() and
    // Evaluate() until it is set again; null or empty: every atom isotropic (the model's own aniso is not read).
    void SetAniso(const std::vector<AtomAniso> *aniso) { aniso_ = aniso; }
    // dT/d(U11 U22 U33 U12 U13 U23) of every anisotropic atom at the last Evaluate() with a gradient, zero for the
    // isotropic ones; empty where none is anisotropic.
    const std::vector<std::array<double, 6>> &UGradient() const { return grad_u_; }

    // 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;
    }

protected:
    bool AnyAniso() const {
        if (aniso_ != nullptr)
            for (const AtomAniso &a : *aniso_)
                if (a.on)
                    return true;
        return false;
    }
    const std::vector<AtomAniso> *aniso_ = nullptr;
    std::vector<std::array<double, 6>> grad_u_;

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 density an atom is gridded out to, in e/A^3 of its blurred density (DensityCalculator::cutoff, 1e-5 by
// default): the radius of every atom, and with it the gather and the gradient, follow from it.
constexpr float REFINEMENT_DENSITY_CUTOFF = 1e-4f;

// 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. An atom with a nonzero aniso is gridded anisotropically, as
// PutModelDensityOnGrid() grids it: its B derivative is that along U + dB / (8 pi^2) I, and with `grad_u` (resized
// to the atoms, zero for the isotropic ones) its derivatives in U11 U22 U33 U12 U13 U23.
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
                          std::vector<std::array<double, 4>> &grad,
                          std::vector<std::array<double, 6>> *grad_u = nullptr);