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

#pragma once

#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). Residues it does not describe are held fixed and
//    stay in the structure factors and the bulk-solvent mask. The weight of the X-ray term against the
//    geometry is chosen from a ladder of full refinements, by R-free among those whose geometry stays within
//    bounds; then whether, and at what weight, to refine the B factors as well, the same way.
//
// 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 group_b = true;             // the placement stage's B shift per residue
    int macro_cycles = 10;
    int xyz_iterations = 10;         // L-BFGS iterations per macro-cycle
    int b_iterations = 5;
    // The positions' X-ray weight, times the gradient-balanced one (BalancedWeight), tried in this order.
    std::vector<double> weight_ladder = {1.0 / 1024, 1.0 / 512, 1.0 / 256, 1.0 / 128, 1.0 / 64, 1.0 / 32};
    // The X-ray weights of the B refinement tried next, times the positions' chosen weight.
    std::vector<double> adp_ladder = {1.0, 4.0, 16.0};
    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
};

// 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 (hydrogens removed, B made isotropic).
    double r_work_start = 1, r_free_start = 1;
    // 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;
    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);