// 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 "gemmi/math.hpp"
#include "gemmi/metadata.hpp"   // TlsGroup
#include "gemmi/model.hpp"

// The anisotropic side of restrained refinement (Refinement.h): TLS groups, and the restraints of atoms refined
// with a U of their own. Atoms are numbered in model order, as everywhere in the refinement.

// A group of atoms whose displacement is in part that of a rigid body: U_TLS(r) = T + A L A^T + A S + S^T A^T at
// r - origin = (x, y, z), A = [[0, z, -y], [-z, 0, x], [y, -x, 0]]. T in A^2, L in rad^2, S in rad A (the files'
// L and S are in degrees). The trace of S does not enter U: S33 is -(S11 + S22).
// Following Schomaker & Trueblood (1968) Acta Cryst. B24, 63-76, and Winn, Isupov & Murshudov (2001) Acta Cryst. D57, 122-133
struct TlsBody {
    std::string id;
    gemmi::Position origin;
    gemmi::SMat33<double> T{0, 0, 0, 0, 0, 0}, L{0, 0, 0, 0, 0, 0};
    gemmi::Mat33 S{0};
    std::vector<int> atoms;   // in model order

    gemmi::SMat33<double> U(const gemmi::Position &pos) const;
    // The 20 parameters (T11 T22 T33 T12 T13 T23, L likewise, S11 S22 S12 S13 S21 S23 S31 S32) and back.
    static constexpr int PARAMETERS = 20;
    std::array<double, PARAMETERS> Get() const;
    void Set(const std::array<double, PARAMETERS> &p);
    // dU/dp at `pos`: per parameter the six U11 U22 U33 U12 U13 U23 (U is linear in the parameters).
    std::array<std::array<double, 6>, PARAMETERS> Derivatives(const gemmi::Position &pos) const;
    // The derivatives of a target in the parameters, from its derivatives in U11 U22 U33 U12 U13 U23 of every atom
    // (`grad_u`, in model order as `atoms`): sum over the body's atoms of grad_u . dU/dp, by the chain rule in matrix
    // form - with G the symmetric matrix of an atom's derivatives (an off-diagonal element half the U_ab one),
    // dT = tr(G dU) gives G for T, A^T G A for L and 2 A^T G for S.
    std::array<double, PARAMETERS> Gradient(const std::vector<gemmi::Atom *> &atoms,
                                            const std::vector<std::array<double, 6>> &grad_u) const;
};

// The TLS groups of a model's metadata (mmCIF _pdbx_refine_tls*, PDB REMARK 3) with their atoms: the atoms of every
// residue of a selection's chain from its first residue to its last, waters left out; a group without numbers or
// without atoms is left out.
std::vector<TlsBody> TlsBodiesFromModel(const gemmi::Model &model, const std::vector<gemmi::TlsGroup> &groups);

// The bodies as the metadata holds them (degrees), each selection one run of consecutive residues of a chain.
std::vector<gemmi::TlsGroup> TlsBodiesToGroups(const gemmi::Model &model, const std::vector<TlsBody> &bodies);

// The restraints of atoms refined with a U of their own (in A^2, U11 U22 U33 U12 U13 U23 per atom):
// - rigid bond, for each pair of bonded or angle-related atoms (the pairs of the B similarity): the difference of
//   their U along the line between them, |(U_i - U_j) n|^2 / sigma^2, sigma 0.004 A^2 - the components U33, U13 and
//   U23 of the difference in a frame whose z is n, as SHELXL's RIGU restrains them;
// - similarity, for the same pairs: |U_i - U_j|^2 (Frobenius) / sigma^2, sigma 0.04 A^2 (SHELXL's SIMU);
// - near isotropy, for an atom in no pair (a water, an ion): |U - tr(U)/3 I|^2 / sigma^2, sigma 0.1 A^2 (ISOR).
// Following Hirshfeld (1976) Acta Cryst. A32, 239-244, and Thorn, Dittrich & Sheldrick (2012) Acta Cryst. A68, 448-451
struct URestraints {
    static constexpr double SIGMA_RIGID = 0.004, SIGMA_SIMILAR = 0.04, SIGMA_ISOTROPIC = 0.1;
    std::vector<std::pair<int, int>> pairs;
    std::vector<int> isotropic;
};
// The pairs of `b_similarity_pairs` whose atoms both refine a U, and the refined atoms of no pair.
URestraints MakeURestraints(const std::vector<std::pair<int, int>> &b_similarity_pairs, const std::vector<char> &refines_u);
// The sum of E at the positions `pos` and the U of every atom `u` (atoms that refine none are not read) and, with
// `grad` (sized as u), dE/dU per atom added to it.
double EvaluateURestraints(const URestraints &r, const std::vector<gemmi::Position> &pos,
                           const std::vector<std::array<double, 6>> &u, std::vector<std::array<double, 6>> *grad);

// Contiguous segments of a chain (TLSMD's idea): residues cut into `k` runs so that the B factors of each run's atoms
// are best fitted, in least squares, by an isotropic TLS model of its own - B(r) a general quadratic in r (tr(A L A^T)
// is tr(L)|r|^2 - r^T L r and tr(A S) linear in r, ten parameters as Painter & Merritt's isotropic TLS), by dynamic
// programming over where the runs end. `residue_atoms`: per residue (in chain order) its atoms, `pos` and `b` per
// atom. Returns per segment its first residue (the first is 0); fewer than k where the chain has too few residues for
// `min_residues` each.
// Following Painter & Merritt (2006) Acta Cryst. D62, 439-450
std::vector<int> TlsSegments(const std::vector<std::vector<int>> &residue_atoms, const std::vector<gemmi::Position> &pos,
                             const std::vector<double> &b, int k, int min_residues);