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

#pragma once

#include <array>
#include <atomic>
#include <cmath>
#include <complex>
#include <memory>
#include <string>
#include <vector>

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

#include "../../common/Reflection.h"   // MergedReflection
#include "DensityBlobs.h"
#include "ModelValidation.h"           // ModelMaps
#include "RefinementRestraints.h"      // GeometryStatistics

class Logger;

// Refinement of an atomic model already placed against the data (ValidateAgainstModel's placed model, or a
// molecular-replacement placement), in two stages:
//
// 1. Placement: each chain as a rigid body, then one B shift per residue, each kept only where it lowers
//    R-free. Needs no dictionary.
// 2. Restrained: every atom's position and isotropic B, against the maximum-likelihood amplitude target with
//    the geometry the dictionary describes (MonomerLibrary.h: bonds, angles, planes, chiral volumes,
//    torsions, repulsions; B similarity of bonded atoms and of atoms in contact, in proportion to B). Residues
//    it does not describe are held fixed and stay in the structure factors and the bulk-solvent mask. The
//    weights of the X-ray term against the geometry and against the B similarity, and whether to refine the B
//    factors, are chosen by R-free among trials whose geometry stays within bounds: a ladder from the placed
//    model, then a trajectory of steps that each try the weights as they are and changed, from where the last
//    step ended, until no step gains; then a short final ladder.
//
// The library reads and writes nothing: the model and the reflections come in, the refined model, the
// statistics and the map coefficients go out. Nothing it finds feeds back into processing decisions.
struct RefinementSettings {
    bool restrained = true;          // stage 2; false stops after the placement stage
    bool score_only = false;         // only score the model as given (RefinementResult::score_start), no stage
    bool group_b = true;             // the placement stage's B shift per residue
    int macro_cycles = 2;            // per weight trial, from the model the trial starts from
    int xyz_iterations = 30;         // L-BFGS iterations per macro-cycle
    int b_iterations = 20;
    // The first trials, from the placed model: the positions' X-ray weight, times the gradient-balanced one
    // (BalancedWeight), each with the B factors' X-ray weight their own balanced one times each of adp_ladder (0: B
    // held).
    std::vector<double> weight_ladder = {1.0 / 1024, 1.0 / 128, 1.0 / 16};
    std::vector<double> adp_ladder = {0.0, 1.0};
    // Then at most this many steps of the trajectory, each a set of trials from where the last one ended,
    int max_steps = 8;
    // and the final trials from where it ended, at its weights times each of these;
    std::vector<double> final_ladder = {0.5, 1.0, 2.0};
    double max_weight = 1.0 / 8;     // the positions' X-ray weight at most, times the balanced one
    bool riding_hydrogens = true;    // the dictionary's hydrogens added, riding on their parents
    bool gpu = true;
    size_t nthreads = 1;
    size_t gpu_resident_bytes = 0;   // held on the card by what runs beside
    // The most weight trials run side by side, 0 for as many as the ladders and the cards take; the result does not
    // depend on it (1 runs them one after another, as a check of that).
    size_t max_trials_at_once = 0;
    // Set by the caller (from any thread) to cut a running refinement short: checked between the minimiser's
    // iterations and before every weight trial. The trials that ran are chosen among as usual, those that did not
    // are left out (RefinementResult::cancelled); cut before the first trial, the model is the placement's.
    const std::atomic<bool> *cancel = nullptr;
};

// R by resolution: one of RefinementScore's shells, equal in width in 1/d^2 between the coarsest and the
// finest reflection.
struct RefinementShell {
    double d_max = 0, d_min = 0;
    int n_work = 0, n_free = 0;
    double r_work = NAN, r_free = NAN;
};

// What scoring a model against the data gives, every term of it: the model's structure factors are
// F_model = k_overall exp(-s^T B s / 4) (Fcalc + k_sol exp(-b_sol s^2 / 4) Fmask), the scale and the solvent
// fitted to the working reflections; R-work and R-free; and the likelihood target, -ln L summed over the
// working and over the free reflections, at D and the variance estimated from the free set.
struct RefinementScore {
    double k_overall = 0;
    double b_iso = 0;                // the overall B: a third of the trace of B (Cartesian, A^2)
    std::array<double, 6> b_aniso{}; // B less b_iso: B11 B22 B33 B12 B13 B23 (Cartesian, A^2)
    double k_sol = 0, b_sol = 0;
    double mask_spacing = 0;         // A, the grid the solvent mask is drawn on
    double r_work = NAN, r_free = NAN;
    int n_work = 0, n_free = 0;
    double nll_work = NAN, nll_free = NAN;
    double fom = NAN;                // the likelihood's mean figure of merit, working reflections
    int hydrogens = 0;               // riding hydrogens in the model scored
    std::vector<RefinementShell> shells;
    // Per reflection, with settings.score_only only: the atoms' Fcalc, the solvent's k_sol exp(-b_sol s^2/4) Fmask
    // and k_overall exp(-s^T B s / 4), in the order of `hkl`, and the likelihood's parameters.
    std::vector<gemmi::Miller> hkl;
    std::vector<std::complex<float>> fcalc, fmask;
    std::vector<float> k_total;
    std::vector<float> d, var;       // the likelihood's D and variance (epsilon beta + sigma^2)
};

// One rung of the weight ladder.
struct RefinementWeightTrial {
    double weight = 0;               // the X-ray term's weight
    double adp_weight = 0;           // the B refinement's X-ray weight; 0: the B factors were held
    double r_work = 1, r_free = 1;
    GeometryStatistics geometry;
    bool within_bounds = false;      // the geometry bounds the choice is made inside
};

struct RefinementResult {
    bool ok = false;
    std::string failure_reason;
    bool on_gpu = false;
    double seconds = 0;
    double d_min = 0;
    int n_work = 0, n_free = 0;

    // The model as given, scaled to these reflections (B made isotropic; the dictionary's hydrogens riding
    // where the restrained stage runs or the model is only scored), and every term of that score.
    double r_work_start = 1, r_free_start = 1;
    RefinementScore score_start;
    // After the placement stage.
    int rigid_bodies = 0;            // the chains moved as rigid bodies, where that was kept
    bool group_b = false;            // the B shift per residue was kept
    double r_work_placement = 1, r_free_placement = 1;
    // After the restrained stage (NaN where it did not run).
    bool restrained = false;
    bool cancelled = false;          // cut short by RefinementSettings::cancel
    double r_work = NAN, r_free = NAN;
    double weight = 0;
    double adp_weight = 0;           // 0: the B factors kept as placed
    std::vector<RefinementWeightTrial> trials;
    GeometryStatistics geometry_start, geometry;

    int atoms = 0, refined_atoms = 0;
    int restrained_residues = 0, fixed_residues = 0;
    std::vector<std::string> fixed_names;
    double k_sol = 0, b_sol = 0;
    double mean_fom = 0;

    std::shared_ptr<gemmi::Structure> model;   // the refined model, in the data's cell and frame
    std::shared_ptr<ModelMaps> maps;           // its 2mFo-DFc and mFo-DFc coefficients (as _maps.mtz holds them)
    // The mFo-DFc map of the refined model and its rms, for finding what it leaves unexplained (DensityBlobs.h).
    std::shared_ptr<gemmi::Grid<float>> difference_map;
    double difference_map_rms = 0;
    // The anomalous difference map on the refined model's phases, where the merge kept the Bijvoet pairs.
    std::shared_ptr<gemmi::Grid<float>> anomalous_map;
    double anomalous_map_rms = 0;
    // What the refined model leaves unexplained in the difference map (FindDensityBlobs).
    std::vector<DensityBlob> blobs;
};

// `placed`: the model in the data's cell and space group. `merged`: the merge, with French-Wilson amplitudes
// and free flags, in the same group.
RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<MergedReflection> &merged,
                             const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A,
                             Logger &logger, const RefinementSettings &settings);