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

#pragma once

// Molecular replacement (rugnux --mr, experimental): a global search for a model's orientation and
// position in the cell, so that a model with no placement of its own - an AlphaFold prediction, or an
// experimental structure from another crystal - can be scored against the data, and the space groups the
// data cannot tell apart can be scored against it.
//
//  - The model: a prediction with residues of pLDDT < 70 removed and B-factors from pLDDT; an experimental
//    structure as its polymer, B-factors kept. Hydrogens dropped either way.
//  - Rotation: a uniform sampling of the point group's rotation ASU, each orientation scored by the
//    correlation of the observed E^2 with the model's own symmetry-summed |F|^2 - read off the
//    model's molecular transform, computed once by FFT and interpolated per orientation.
//  - Translation: the full-symmetry translation function on one FFT per orientation and space group.
//  - Further copies, one at a time, as many as the cell can hold: the translation function again with
//    the copies already placed as a known partial structure, each new copy refined against the others.
//  - Rescoring: the best placements through the rigid body, bulk solvent, R-free, a packing check
//    against the symmetry mates and a null of the same model in random orientations.
//
// The rotation and translation functions, the rigid body and the structure factors the R-factors are
// taken on run on the GPU where there is one, and on the CPU otherwise; the same functions either way.

#include <complex>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "gemmi/model.hpp"
#include "gemmi/recgrid.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

#include "../../common/Reflection.h"   // MergedReflection
#include "../../common/UnitCell.h"

class Logger;

// The model as the search uses it, centroid at the origin. A prediction (AlphaFold, ModelCIF) trimmed to the
// residues it is confident of, B-factors from pLDDT; an experimental structure as it is, its polymer only.
struct SearchModel {
    gemmi::Structure st;
    bool predicted = true;     // a prediction, its B column a pLDDT; an experimental structure otherwise
    int residues_read = 0;
    int residues_kept = 0;
    double mass_da = 0;        // of the residues kept, 110 Da each
    double radius_A = 0;       // the farthest atom from the centroid
    double error_A = 0;        // rms of the coordinate errors the pLDDTs predict, over the atoms kept (fixed for an
                               // experimental structure)
    std::vector<int> residue_index;   // of each residue kept, in model order: its position in the prediction
    std::string description;   // "the whole model", or which domain
    std::string failure_reason;
};
// A prediction where the file says so - ModelCIF's _ma_qa_metric - or where it has no cell and its B column
// holds one value per residue, from 0 to 100; an experimental structure otherwise.
SearchModel PrepareSearchModel(const std::string &path);

// The rigid domains of a prediction from its predicted aligned error (PAE, the AlphaFold DB's JSON: a matrix
// "predicted_aligned_error" over the residues of the whole prediction): the kept residues clustered by average
// linkage on the symmetrised PAE until every two clusters are a mean PAE_DOMAIN_A apart; clusters of at least
// DOMAIN_MIN_RESIDUES, largest first, each a list of positions in the prediction. Empty where the file cannot
// be read or does not match the model.
std::vector<std::vector<int>> PaeDomains(const std::string &pae_path, const SearchModel &model);
// The search model of one domain: the residues of `model` at those positions, recentred.
SearchModel DomainModel(const SearchModel &model, const std::vector<int> &residues);
// The PAE file the AlphaFold DB puts beside a model: AF-<id>-F1-model_v<n>.cif ->
// AF-<id>-F1-predicted_aligned_error_v<n>.json in the same directory; empty where the name does not follow it.
std::string PaePathFor(const std::string &model_path);

// Rotations covering SO(3) with neighbours about `step_deg` apart: the identity, then shells of
// rotation angle kappa = step, 2 step, ... 180 deg, each with its axes on a golden-spiral lattice of
// the sphere as dense as 2 sin(kappa/2) - the distance a change of axis moves a rotation - needs.
std::vector<gemmi::Mat33> UniformRotations(double step_deg);

// The rotations of the space group's point group, in Cartesian coordinates of `cell`.
std::vector<gemmi::Mat33> CartesianPointGroup(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell);

// The rotations of `all` in the rotation ASU of `point_group`: the crystal puts a copy of the model
// in every orientation S R, so only one rotation of each class {S R} is worth scoring - here the one
// closest to the identity (largest trace), the classes' Dirichlet region about it.
std::vector<gemmi::Mat33> RotationAsu(const std::vector<gemmi::Mat33> &all,
                                      const std::vector<gemmi::Mat33> &point_group);

// The angle in degrees between a and b, the nearest over the classes {S b}.
double RotationDistanceDeg(const gemmi::Mat33 &a, const gemmi::Mat33 &b,
                           const std::vector<gemmi::Mat33> &point_group);

// The structure factors of one copy of the model, at its centroid, at any reciprocal vector.
class MolecularTransform {
public:
    MolecularTransform(const SearchModel &model, double d_min, size_t nthreads);
    // F at the Cartesian reciprocal vector s (1/A), interpolated trilinearly.
    std::complex<double> At(const gemmi::Vec3 &s) const;
    // The grid At() interpolates (half-l, XYZ), its box edge in A and the blur it is unblurred by.
    const gemmi::FPhiGrid<float> &Grid() const { return grid_; }
    double Box() const { return box_; }
    double Blur() const { return blur_; }
    double DMin() const { return d_min_; }

private:
    gemmi::FPhiGrid<float> grid_;   // half-l, on a P1 box `box_` A wide
    double box_ = 0;
    double blur_ = 0;
    double d_min_ = 0;
};

// The observed reflections the search scores against: E^2 in resolution bins, with each unique
// reflection's multiplicity on the whole sphere.
struct SearchReflections {
    std::vector<gemmi::Miller> hkl;
    std::vector<double> e2;
    std::vector<double> multiplicity;
    std::vector<int> bin;
    int bins = 0;
};
SearchReflections MakeSearchReflections(const std::vector<MergedReflection> &merged, const gemmi::UnitCell &cell,
                                        const gemmi::SpaceGroup &sg, double d_min, double d_max);

// Copy j of the model contributes |F_model(s)|^2 at s = h R_j, the reflection's index under the j-th
// operator of the point group as a Cartesian reciprocal vector: refl.hkl.size() x |point group|, index-major.
std::vector<gemmi::Vec3> CopyVectors(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                     const gemmi::SpaceGroup &sg);

// The rotation function at each of `rotations`: the correlation of E^2 - 1 with the bin-normalised
// sum over the point group's copies of |F_model|^2.
std::vector<double> RotationFunction(const MolecularTransform &ft, const SearchReflections &refl,
                                     const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                     const std::vector<gemmi::Mat33> &rotations, size_t nthreads);

// What the translation function of one space group needs that does not depend on the orientation:
// for every reflection not absent in the group (a "term") and every operator j, s = h R_j as a
// Cartesian vector and exp(2 pi i h.t_j); the term's weight - its multiplicity times E^2 - 1 - and bin;
// and for every term and every ordered pair of operators (j, k), j != k, where the coefficient goes:
// the index h R_j - h R_k on the function's grid, or u = -1 where its l is negative (only l >= 0 is kept).
struct TranslationSetup {
    int nu = 0, nv = 0, nw = 0;      // the function's grid
    int ops = 0;
    std::vector<gemmi::Vec3> s;                     // terms x ops
    std::vector<std::complex<double>> phase;        // terms x ops
    std::vector<double> weight;                     // terms
    std::vector<double> e2, multiplicity;           // terms
    std::vector<int> bin;                           // terms
    int bins = 0;
    std::vector<int> pair_j, pair_k;                // the ordered pairs
    std::vector<int> target_u, target_v, target_l;  // terms x pairs
    // With a known partial structure (MakeTranslationSetup's `known_part`): where the cross terms
    // conj(F_known) g_j go, at h R_j - or, as their conjugate F_known conj(g_j), at -h R_j where that
    // has l >= 0; both where l = 0. One entry per coefficient: its term, operator, index and whether it
    // is the conjugate.
    std::vector<int> cross_term, cross_op, cross_u, cross_v, cross_l;
    std::vector<char> cross_conj;
    size_t Terms() const { return weight.size(); }
};
// `known_part`: the grid is sized for the cross terms with a known partial structure as well, which
// reach out to h R_j where the copies' own terms reach h R_j - h R_k (in a polar group no further than
// the plane the polar axis is free in).
TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                      const gemmi::SpaceGroup &sg, bool known_part = false);

// F of copies of the model already placed - copy c turned by rotations[c] about its centroid, the
// centroid at fractional positions[c] - at each term of `setup`, summed over the copies and their
// symmetry mates, in the units the translation function's g_j are in.
std::vector<std::complex<double>> KnownStructureFactors(const MolecularTransform &ft, const TranslationSetup &setup,
                                                        const gemmi::UnitCell &cell,
                                                        const std::vector<gemmi::Mat33> &rotations,
                                                        const std::vector<gemmi::Fractional> &positions,
                                                        size_t nthreads = 1);

// The translation function of the model in orientation `rotation`: its highest point as a Z-score over
// every point of the function, and where it is (the fractional position of the model's centroid).
struct TranslationPeak {
    double z = 0;
    gemmi::Fractional t;
};
// With `known` (KnownStructureFactors on a setup made with known_part), the function of one more copy
// beside the known ones: the overlap of E^2 - 1 with |F_known + F(t)|^2, whose cross term
// 2 Re(conj(F_known) F(t)) is a Fourier series in t too and goes on the same grid.
TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
                                    const gemmi::Mat33 &rotation,
                                    const std::vector<std::complex<double>> *known = nullptr);

// --- The likelihood a placement is judged by ---
// The terms it is summed over: the reflections of a SearchReflections not absent in a group, each with
// its E_obs, whether it is centric, its epsilon and 1/d^2, and for every operator j of the group h R_j (as
// a Cartesian vector and as indices) and exp(2 pi i h.t_j), as in TranslationSetup.
struct LikelihoodSetup {
    int ops = 0;
    std::vector<gemmi::Vec3> s;                     // terms x ops
    std::vector<gemmi::Miller> k;                   // terms x ops
    std::vector<std::complex<double>> phase;        // terms x ops
    std::vector<double> eo, epsilon, s2;            // terms
    std::vector<char> centric;                      // terms
    std::vector<int> bin;                           // terms
    int bins = 0;
    int max_k[3] = {0, 0, 0};                       // the largest |index| of h R_j along each axis
    size_t Terms() const { return eo.size(); }
};
LikelihoodSetup MakeLikelihoodSetup(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                    const gemmi::SpaceGroup &sg);

// g_j = G_j(h) exp(2 pi i h.t_j) of one copy in orientation `rotation`, terms x ops: the copy with its
// centroid at t adds sum_j g_j exp(2 pi i (h R_j).t) to F(h). G_j(h) = F1(h R_j), F1 the transform of
// the copy in `cell` with its centroid at the origin - one FFT of the cell's grid to `d_min`, read off at
// the lattice points themselves, so exact at any resolution. Where `ft` is given and reaches `d_min`, read
// off it instead, interpolated as the rotation and translation functions read it.
std::vector<std::complex<double>> CopyFactors(const SearchModel &model, const LikelihoodSetup &setup,
                                              const gemmi::UnitCell &cell, const gemmi::Mat33 &rotation,
                                              double d_min, size_t nthreads, const MolecularTransform *ft = nullptr);
// `f` += that copy's F with its centroid at the fractional position t.
void AddCopy(const LikelihoodSetup &setup, const std::vector<std::complex<double>> &g, const gemmi::Fractional &t,
             std::vector<std::complex<double>> &f);

// The log-likelihood gain of the model structure factors `f` (one per term) over a random structure,
// with sigma_A = scale exp(-2 pi^2 error^2 / (3 d^2)): E_calc are |f| normalised in the resolution bins.
double LogLikelihoodGain(const LikelihoodSetup &setup, const std::vector<std::complex<double>> &f, double error_A,
                         double scale);
// The same with the scale refined to its maximum (returned in `scale`).
double RefinedLogLikelihoodGain(const LikelihoodSetup &setup, const std::vector<std::complex<double>> &f,
                                double error_A, double *scale = nullptr);

// How the likelihood of one more copy in orientation `rotation` at `t`, beside the `known` structure factors
// (zero for none), stands out over the same copy in the same orientation at random positions: the gain
// with the scale refined at t, and its Z-score over the gains at those positions with that scale held.
struct LikelihoodPeak {
    double llg = 0;       // of all the copies, the known and this one
    double gain = 0;      // over the known ones alone, at their own refined scale
    double z = 0;
    double scale = 0;
};
// `g`: CopyFactors of the copy in that orientation; `known_llg`: RefinedLogLikelihoodGain of `known` alone.
LikelihoodPeak LikelihoodTranslationZ(const LikelihoodSetup &setup, const std::vector<std::complex<double>> &known,
                                      double known_llg, const std::vector<std::complex<double>> &g,
                                      const gemmi::Fractional &t, double error_A, size_t nthreads);
// The gain alone: of the known structure factors and the copy over the known ones (`llg`: of all of them).
double LikelihoodGain(const LikelihoodSetup &setup, const std::vector<std::complex<double>> &known, double known_llg,
                      const std::vector<std::complex<double>> &g, const gemmi::Fractional &t, double error_A,
                      double *llg = nullptr, double *scale = nullptr);

// The model placed: rotated about its centroid by `rotation`, the centroid at `t`, in `cell` and `sg`.
gemmi::Structure PlaceSearchModel(const SearchModel &model, const gemmi::Mat33 &rotation, const gemmi::Fractional &t,
                                  const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg);

// The fraction of the C-alpha atoms within 3 A of a C-alpha of another copy: a symmetry mate, or one of
// the other `copies` the structure's chains are split into evenly (in order), or their mates.
double ClashFraction(const gemmi::Structure &placed, size_t copies = 1);

// Copies of the model as one structure: their chains in order, renamed so that no two share a name.
gemmi::Structure AssembleCopies(const std::vector<gemmi::Structure> &copies);

// The copies of an assembly (AssembleCopies), its chains split evenly in order.
std::vector<gemmi::Structure> SplitCopies(const gemmi::Structure &assembly, int copies);

// The setting of a subgroup's merge (cell `h`, group `h_sg`) in the cell `g` of a group `g_sg`: `basis` has h's
// axes as columns, lattice vectors of g (with g's centring) in g's fractional coordinates, right-handed and
// matching h's lengths and angles to 1%; `origin` is where g's origin lies in h's fractional coordinates, so
// that the operators of g with h's rotations are h's own: x_h = basis^-1 x_g + origin. Empty where there is no
// such setting.
struct SubgroupSetting {
    gemmi::Mat33 basis;
    gemmi::Vec3 origin;
};
std::optional<SubgroupSetting> SubgroupBasis(const gemmi::UnitCell &g, const gemmi::SpaceGroup &g_sg,
                                             const gemmi::UnitCell &h, const gemmi::SpaceGroup &h_sg);

// The `copies` of a placement in group `g_sg` and every symmetry image of them that the subgroup `h_sg` does
// not make itself, as independent copies in the subgroup's cell `h`, in its setting (SubgroupBasis); their
// number in `copies_out`.
gemmi::Structure ExpandIntoSubgroup(const gemmi::Structure &placed, int copies, const gemmi::SpaceGroup &g_sg,
                                    const SubgroupSetting &setting, const gemmi::UnitCell &h,
                                    const gemmi::SpaceGroup &h_sg, int *copies_out);

// The probabilities of 1, 2, ... copies of a model of `mass_da` in the asymmetric unit of a cell of
// `asu_volume` A^3, from the distribution of the Matthews coefficient over protein crystals; one entry
// per copy number whose Vm is at least 1.5 A^3/Da.
std::vector<double> MatthewsProbabilities(double asu_volume, double mass_da);

// The space groups with the rotations and the centring of `sg`, in its setting: those merged
// intensities cannot tell apart, apart from systematic absences. `sg` itself first.
std::vector<gemmi::SpaceGroup> SameRotationGroups(const gemmi::SpaceGroup &sg);

// Translational non-crystallographic symmetry: copies in the same orientation, a translation apart, make
// the native Patterson peak off its origin where no lattice vector is. The highest peak of the Patterson
// of E^2 (5-10 A), at least 15 A from the origin and from every centring vector, as a fraction of the
// origin's height, and where it is; tNCS where that is a fifth or more.
struct TranslationalNcs {
    bool present = false;
    double height = 0;              // over the origin's
    gemmi::Fractional vector;
};
TranslationalNcs DetectTranslationalNcs(const std::vector<MergedReflection> &merged, const gemmi::UnitCell &cell,
                                        const gemmi::SpaceGroup &sg);

// Whether groups `a` and `b` (the same rotations) differ only in screw translations that the tNCS vector
// `v` makes up: then copies a translation v apart put one group's operators onto the other's, and neither
// is told from the other by a search that has not placed all of them.
bool TncsRelatesGroups(const gemmi::SpaceGroup &a, const gemmi::SpaceGroup &b, const gemmi::Fractional &v);

// One merge the model is searched against: the answer's, or the twin hypothesis' in a subgroup.
struct MRMerge {
    std::string name;                                // "MERGE" or "TWIN_SUBGROUP_MERGE"
    const std::vector<MergedReflection> *merged = nullptr;
    UnitCell cell;
    gemmi::SpaceGroup space_group;                   // the group it was merged in
    std::optional<gemmi::Op> twin_law;               // the twin hypothesis' merge: the law it carries
};

// One space group's best placement: its first copy and every further one the search added.
struct MRCandidate {
    int merge = 0;                                   // index into the merges searched
    gemmi::SpaceGroup space_group;
    double vm = 0;                                   // Matthews coefficient of one copy
    int copies_likely = 1;                           // the most probable copy number, by Matthews
    int copies_max = 1;                              // the most the search tried (Matthews probability >= 1%)
    int copies = 0;                                  // placed
    double rf_z = 0;                                 // of the first copy's orientation, over the rotation function
    double tf_z = 0;                                 // the highest of the copies'
    std::vector<double> copy_tf_z;                   // each copy's, the later ones beside the earlier
    // The likelihood (LikelihoodTranslationZ), after the rigid body, taken only for a copy whose translation
    // function does not stand out (NaN otherwise): of all the copies with the last such one, and the highest
    // of their Z-scores, and each copy's.
    double llg = NAN;
    double llg_z = NAN;
    std::vector<double> copy_llg_z;
    double clash_fraction = 0;                       // of all the copies
    double r_work = NAN, r_free = NAN;               // after the rigid body, to d_min_score
    double k_sol = NAN, b_sol = NAN;                 // the bulk solvent R was computed with
    double null_sigma = NAN;                         // R-work above its random-orientation null, in sd
    // On the twin hypothesis' merge, against the data as twinned by its law (the twin comparison).
    double r_work_twin = NAN, r_free_twin = NAN, twin_fraction = NAN;
    std::string verdict;                             // PLACED, or why not
    std::shared_ptr<gemmi::Structure> placed;        // every copy
};

struct MolecularReplacementResult {
    bool ok = false;
    std::string failure_reason;
    std::string model_path;
    int residues_read = 0, residues_kept = 0;
    double mass_da = 0;
    std::string search_model;                        // the whole model, or the PAE domain the decision was made with
    bool predicted_model = true;                     // SearchModel::predicted
    int pae_domains = 0;                             // the rigid domains of the prediction's PAE; 0 without a PAE file
    double d_min_search = 0, d_max_search = 0;
    double d_min_score = 0;                          // the limit R-work / R-free are computed to
    std::vector<MRCandidate> candidates;
    // The decision: the candidate the model chooses, or -1 where it chooses none, and in words why.
    int chosen = -1;
    std::string decision;
    bool gpu_search = false;   // the rotation and translation functions ran on the GPU
    TranslationalNcs tncs;     // of the answer's merge
    // An ENANTIOMORPH decision is taken into the answer's label only as model validation takes a model's
    // enantiomorph (ValidateAgainstModel): the placed model has to fit against its null and the anomalous
    // map at its atoms must not be inverted. Filled by the caller, which runs that validation.
    struct HandCheck {
        bool performed = false;
        bool adopted = false;
        bool model_fits = false;
        double fit_sigma = NAN;
        bool anomalous_inverted = false;
        int anomalous_pairs = 0;
        double anomalous_scatterer_mean_sigma = NAN;
        std::string deepest_site;
        double deepest_sigma = NAN;
        double seconds = 0;
    } hand_check;
    // Seconds per step.
    double t_model = 0, t_transform = 0, t_rotation = 0, t_translation = 0, t_rescore = 0, t_null = 0;
    double t_domains = 0;   // the PAE domains' own searches, wall clock (they run side by side)
};

// Search `merges` (the answer's first) for the predicted model at `model_path`, in every space group with
// the rotations of each merge's group - one copy, then further copies one at a time while each new one
// stands out, packs and lowers R-free - and decide. A candidate is PLACED where a copy's translation
// function stands out (TF Z >= 8), the copies pack, one copy fits the cell, and its R-work beats its own
// null. The model chooses among the answer merge's groups where one placed candidate leads
// every other by an R-free margin; it chooses the twin hypothesis' merge only where none of the
// answer's groups can hold the model at all (Matthews, packing) and a group of the subgroup can.
MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector<MRMerge> &merges,
                                                   Logger &logger, size_t nthreads);

// The hand check of an ENANTIOMORPH decision, filled into result.hand_check: taken into the answer by
// --model's rule for a model's enantiomorph - the model has to fit against its own null, which the choice
// already required of it, and the anomalous map at its atoms (`merged`, the answer's merge in
// `data_space_group`) must not come out inverted, read by model validation's own code on the placement
// and the bulk solvent MR scored. Where it is not taken, the decision becomes ENANTIOMORPH_VETOED.
void CheckEnantiomorphHand(MolecularReplacementResult &result, const std::vector<MergedReflection> &merged,
                           const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads);