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

#pragma once

#include <array>
#include <cmath>
#include <string>
#include <vector>

#include <memory>

#include "gemmi/model.hpp"
#include "gemmi/monlib.hpp"
#include "gemmi/topo.hpp"
#include "gemmi/unitcell.hpp"

#include "RefinementMath.h"   // RestraintRecord

// The geometry side of restrained refinement (Refinement.h): what the dictionary says each refinable atom's
// surroundings should look like, as flat lists of restraints the CPU and the GPU evaluate alike
// (RefinementMath.h). Atoms are numbered in model order, as ModelPositions() lists them.

// A list of restraints and how their gradients reach the atoms: each restraint writes one gradient slot per
// atom it involves, and atom a sums the slots slot_index[slot_start[a] ... slot_start[a + 1]) in that order -
// a gather, so that a GPU, summing in the same order, needs no atomics and gives the same bits every run.
struct RestraintSet {
    std::vector<RestraintRecord> records;
    std::vector<int> plane_atoms;
    int n_slots = 0;
    std::vector<int> slot_start;   // n_atoms + 1
    std::vector<int> slot_index;

    void Add(RestraintRecord r, int n_atoms_of_r);
    void Append(const RestraintSet &other);   // other's records after these; Finalize() again afterwards
    void Finalize(size_t n_atoms);
    // Which atom each slot belongs to (Finalize() works it out); -1 for none.
    std::vector<int> slot_atom;
};

// The value of one restraint's own quantity (a distance, an angle, a volume ...) at `xyzb`, for statistics.
double RestraintValue(const RestraintSet &set, const RestraintRecord &r, const std::vector<std::array<double, 4>> &xyzb);

// The sum of E over the set at `xyzb` (x, y, z, B per atom) and, with `grad`, dE/d(x, y, z, B) per atom (added
// to it). On the CPU, in double.
double EvaluateRestraints(const RestraintSet &set, const std::vector<std::array<double, 4>> &xyzb,
                          std::vector<std::array<double, 4>> *grad);

// The diagonal of the Hessian of the set's B similarity restraints with respect to each atom's B, at `xyzb`.
std::vector<double> BRestraintCurvature(const RestraintSet &set, const std::vector<std::array<double, 4>> &xyzb);

// What refinement knows of a model's chemistry.
struct RefinementTopology {
    RestraintSet geometry;   // bonds, angles, torsions, chiral volumes, planes
    RestraintSet metal_bonds;   // the bonds of metals (ions, a heme's iron) to their donors, apart from the rest
    RestraintSet adp;        // B of bonded and angle-related atoms alike
    std::vector<char> refinable;           // per atom: its residue is described by the dictionary
    // Per atom: the atoms the repulsion leaves out - itself and those up to three bonds away (sorted).
    std::vector<std::vector<int>> excluded;
    std::vector<float> vdw_radius;         // A
    std::vector<char> hbond;               // donor, acceptor or both (water included)
    std::vector<char> metal;               // a metal atom: an ion, or a metal within a ligand (a heme's iron)
    std::vector<char> altloc;
    std::vector<gemmi::El> element;
    // Per atom: the heavy atom a hydrogen rides on, -1 for every other atom.
    std::vector<int> rider_parent;
    // The topology itself and the dictionary it points into, for placing the hydrogens again.
    std::shared_ptr<gemmi::MonLib> monlib;
    std::shared_ptr<gemmi::Topo> topo;
    int restrained_residues = 0;
    int fixed_residues = 0;
    std::vector<std::string> fixed_names;  // the residue names held fixed, each once
    int bonds = 0, angles = 0, torsions = 0, chirals = 0, planes = 0;
    int metal_links = 0;                   // metal-donor bonds made from the model's LINK records (AddMetalLinks)
};

// The topology of the first model of `st` (hydrogens already removed) from the embedded monomer library
// (MonomerLibrary.h), through gemmi's prepare_topology(): links between residues (peptide, cis-peptide,
// nucleic-acid backbone, disulfides, the model's LINK records the library describes) and their
// modifications included; the LINK records of metals (ions, a heme's iron) the library has no link for are
// restrained at the library's metal-site distances (RefinementRestraints.cpp, AddMetalLinks). A residue the library
// does not describe, or with an atom it does not know, is held fixed: it stays in the structure factors and the
// bulk-solvent mask, and it is restrained against.
// Torsions: those with an esd, except the main-chain phi/psi and the sugar puckers given as alternatives.
// With `hydrogens`, the dictionary's hydrogens are added to the residues it describes (waters excepted) and
// placed; they ride on their parent atoms and take part in no restraint.
RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens = false);

// The riding hydrogens placed again on their parents as they are now (gemmi's place_hydrogens_on_all_atoms).
void PlaceRidingHydrogens(RefinementTopology &topo);

// The repulsions of the model as it is now: every pair of atoms, symmetry mates included, closer than its
// contact distance plus `margin` and not excluded, where at least one of the two is refinable and the two are
// not alternative conformations of each other. Contact distances: van der Waals radii summed, 2.6 A where
// both can hydrogen-bond, 1.9 A between a metal and N, O or S. `model` is not changed (gemmi's NeighborSearch
// takes it non-const); the search runs on `nthreads` threads.
RestraintSet BuildRepulsions(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell,
                             const gemmi::SpaceGroup &sg, double margin, size_t nthreads = 1);

// The B similarity of the atoms in contact (closer than 4 A, symmetry mates included) that no bond or angle relates,
// waters among them, under the same conditions as BuildRepulsions().
RestraintSet BuildBContacts(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell,
                            const gemmi::SpaceGroup &sg, size_t nthreads = 1);

// The usual summary of a model's geometry.
struct GeometryStatistics {
    double rmsd_bonds = 0, rmsd_angles = 0, rmsd_chirals = 0;
    double rmsz_bonds = 0, rmsz_angles = 0;   // rms of (deviation / the dictionary's sigma)
    int bonds = 0, angles = 0;
    int clashes = 0;     // repulsions overlapping by more than 0.4 A
    // The metal-donor bonds, apart: their sigma is the spread of real sites (0.04-0.1 A), so they sit further
    // from ideal than covalent bonds do and are not what the bond rmsd's bound is about.
    double rmsd_metal = 0, rmsz_metal = 0;
    int metal_bonds = 0;
    // Bonds and metal bonds together, as one rmsd (what REFMAC lists as its bond distances).
    double RmsdAllBonds() const {
        const int n = bonds + metal_bonds;
        return n > 0 ? std::sqrt((rmsd_bonds * rmsd_bonds * bonds + rmsd_metal * rmsd_metal * metal_bonds) / n) : 0;
    }
};
// `metal_bonds`: where given, its bonds are measured too (rmsd_metal), apart from `geometry`'s.
GeometryStatistics MeasureGeometry(const RestraintSet &geometry, const RestraintSet &repulsions,
                                   const std::vector<std::array<double, 4>> &xyzb,
                                   const RestraintSet *metal_bonds = nullptr);