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

#pragma once

#include <array>
#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);

// What refinement knows of a model's chemistry.
struct RefinementTopology {
    RestraintSet geometry;   // bonds, angles, torsions, chiral volumes, planes
    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 ion
    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;
};

// 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. 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.
RestraintSet BuildRepulsions(const RefinementTopology &topo, const gemmi::Model &model, const gemmi::UnitCell &cell,
                             const gemmi::SpaceGroup &sg, double margin);

// The usual summary of a model's geometry.
struct GeometryStatistics {
    double rmsd_bonds = 0, rmsd_angles = 0, rmsd_chirals = 0;
    int bonds = 0, angles = 0;
    int clashes = 0;     // repulsions overlapping by more than 0.4 A
};
GeometryStatistics MeasureGeometry(const RestraintSet &geometry, const RestraintSet &repulsions,
                                   const std::vector<std::array<double, 4>> &xyzb);