// 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 <set>

#include "gemmi/model.hpp"
#include "gemmi/monlib.hpp"
#include "gemmi/neighbor.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);

// An atom on a crystallographic special position: one that some symmetry operations of the group, each with a
// lattice translation, take onto itself - the site's symmetry. Its position moves only within the subspace those
// operations leave as it is (a 2-fold axis: along the axis; a 4-fold's: along it; a centre of inversion: not at all),
// and its occupancy is the model's, which counts each of the coinciding copies (0.25 on a 4-fold axis).
// The site symmetry and its projection following Grosse-Kunstleve & Adams (2002) Acta Cryst. A58, 60-65
struct SpecialPosition {
    int atom = 0;                           // in model order
    int multiplicity = 1;                   // the site symmetry's order: how many copies coincide, itself included
    std::vector<gemmi::Transform> ops;      // the site symmetry, Cartesian x -> R x + t, the identity left out
    gemmi::Mat33 projector;                 // the mean of the R: projects a shift onto the subspace

    // The point of the subspace nearest to `p`: the mean of p and its images under the site symmetry.
    gemmi::Position OnSite(const gemmi::Position &p) const;
};

// Atoms nearer than this to one of their own symmetry images are on that special position: the depositions of the
// refine and open arms put theirs within 0.03 A of the site, one disordered water at 0.45 A.
constexpr double SPECIAL_POSITION_TOLERANCE = 0.5;   // A

// The site symmetry of the point `p` in `cell` (its images set, UnitCell::set_cell_images_from_spacegroup): no ops,
// multiplicity 1, where it is on no special position.
SpecialPosition SiteSymmetry(const gemmi::UnitCell &cell, const gemmi::Position &p);

// The atoms of the first model of `st` on special positions of its cell and space group, in model order; none
// where the structure has no space group.
std::vector<SpecialPosition> FindSpecialPositions(const gemmi::Structure &st);

// 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)
    int metal_links_refused = 0;           // LINK records at an outer-sphere distance, left as contacts
    // The atoms on special positions (FindSpecialPositions), placed exactly on them; per atom its entry, -1 none.
    std::vector<SpecialPosition> special;
    std::vector<int> special_index;
};

// 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.
// The atoms on special positions are moved onto them (SpecialPosition::OnSite).
RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens = false);

// A metal's inner sphere ends this far beyond the library's distance: Mg-O 2.07 A + 0.33 = 2.4 A, the inner-sphere
// bound of Leonarski, D'Ascenzo & Auffinger (2017) Nucleic Acids Res. 45, 987-1004, taken for every metal.
constexpr double INNER_SHELL_MARGIN = 0.33;           // A

// Atom `i` (in model order; `a` in residue `res`) retyped where its residue was relabelled (a water become an ion, an
// ion a water): its element, van der Waals radius, hydrogen-bond and metal flags taken again, its bonds removed with
// the repulsions they excluded, and - a metal - a bond made to each of `donors` (atom indices), at the library's
// distance for the coordination they make. Its B similarities go with DropBRestraints(), once for all the atoms
// retyped. The atoms stay as they are numbered, so an engine made for them goes on.
void RetypeAtom(RefinementTopology &topo, const gemmi::Residue &res, const gemmi::Atom &a, int i,
                const std::vector<int> &donors);
// The B similarities of the atoms marked in `drop` (one per atom) taken out of topo.adp.
void DropBRestraints(RefinementTopology &topo, const std::vector<char> &drop);


// 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. An atom on a special position (topo.special) makes no contact with
// its own coinciding copies, and each of its contacts is listed once, not once per coinciding copy. 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);

// One atom of a model near a point (AtomsNear): which atom, and its copy - symmetry mate or lattice translate -
// nearest the point.
struct ModelNeighbour {
    int chain = 0, residue = 0, atom = 0;   // indices into the model
    gemmi::Position pos;                    // the copy's position
    double distance = 0;                    // A, from the point
    gemmi::Transform image;                 // x -> R x + t (Cartesian) taking the atom to its copy
};

// The atoms of `model` within `radius` of `point`, symmetry mates and lattice translates included, nearest first;
// the point's own atom, where it is one, among them at distance 0. The coinciding copies of an atom on a special
// position are listed once. `ns` is a search populated over `model` in a
// cell with its symmetry images set (UnitCell::set_cell_images_from_spacegroup), made with a radius of at least
// `radius`; gemmi's NeighborSearch takes it non-const.
std::vector<ModelNeighbour> AtomsNear(gemmi::NeighborSearch &ns, const gemmi::Model &model, const gemmi::Position &point,
                                      double radius);

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