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

#pragma once

#include <array>
#include <complex>
#include <cstddef>
#include <optional>
#include <vector>

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

class Logger;

struct RigidBodyRefineResult {
    bool converged = false;
    double angle_deg = 0.0;          // magnitude of the rotation about the model centroid
    gemmi::Mat33 rotation;           // that rotation, Cartesian (identity where nothing moved)
    double shift_A = 0.0;            // magnitude of the translation
    std::vector<double> zones;       // the resolution ladder actually walked, coarsest first
    int evaluations = 0;             // structure-factor evaluations (each one re-fits the scale)
    int jacobians = 0;               // Jacobians, each three single-copy transforms (not counted above)
    double seconds = 0.0;
    double k_sol = 0.0, b_sol = 0.0; // the bulk solvent the last zone's target was scored through
};

// Every atom position of `model`, in the order the model iterates them.
std::vector<gemmi::Position> ModelPositions(const gemmi::Model &model);
void SetModelPositions(gemmi::Model &model, const std::vector<gemmi::Position> &pos);

// How far, in degrees, a rotation about the centroid can be from the answer and still be walked onto
// it by RefineRigidBody: the angle that moves the model's rms-radius atom by the resolution of the
// ladder's first zone. Inside it the first zone sees the rotated model overlap the density it belongs
// to, and its single broad minimum is that placement.
double RigidBodyReachDeg(const gemmi::Model &model, double d_min);

// Refine the placement of `model` as one rigid body against the observed amplitudes: an angle-axis
// rotation about the model's own centroid followed by a translation, six parameters, over a
// coarse-to-fine resolution ladder. Each evaluation recomputes Fcalc and the bulk-solvent mask for
// the moved model and re-fits the scale (k_overall and the anisotropic B), so the target measures
// the placement and not the scale. The two flat-solvent constants are fitted once per zone, inside
// their physical range, and then held: they describe the crystal rather than the placement.
//
// `fobs` should be the working set only - the free reflections are what the caller decides on. The
// model is left MOVED (whether or not the refinement helped): the caller scores it and puts it back
// with SetModelPositions() if it did not.
//
// `gpu`: evaluate the target on an engine of this pool (RigidBodyGPU.h, CUDA builds only) rather than
// on the CPU. The same target either way; null is the CPU.
//
// `fixed`: other atoms of the crystal that do not move - the other copies of a molecular-replacement
// search - whose Fcalc is added to the body's at every evaluation, and which the bulk-solvent mask covers
// as well; that mask is computed once per zone, with the body where the zone starts, and held. With them
// the origin is no longer free, and the translation is refined in all three directions.
class RigidBodyGPUPool;
RigidBodyRefineResult RefineRigidBody(gemmi::Model &model,
                                      const gemmi::UnitCell &cell,
                                      const gemmi::SpaceGroup &sg,
                                      const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs,
                                      double d_min,
                                      Logger &logger,
                                      size_t nthreads = 1,
                                      RigidBodyGPUPool *gpu = nullptr,
                                      const gemmi::Model *fixed = nullptr);

// --- The pieces RefineRigidBody is built from, exposed for the tests and the GPU backend ---

// The resolution zones RefineRigidBody walks down to d_min, coarsest first, before a zone too thin to
// run in is dropped.
std::vector<double> RigidBodyLadder(double d_min);

// A zone's empty grid: its size, cell and group, as DensityCalculator sizes its own for the zone.
gemmi::Grid<float> RigidBodyZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min);

// The step of the forward-difference rotation columns of a zone's Jacobian, in angstroms of q.
double RigidBodyJacobianStep(double d_min);

// The projector onto the directions in which the origin of `sg` is free (see RefineRigidBody).
gemmi::Mat33 RigidBodyGaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell);

// Fcalc of the whole crystal composed from the transform of ONE copy of its content, gridded on the
// zone's grid without symmetrization: F(h) = n_cen * sum over the group's operators x -> Rx + t of
// exp(+2 pi i h.t) F1(hR). gemmi's symmetrize() maps grid points onto grid points, so this is the
// transform of the symmetrized grid, rearranged - the same number up to rounding - without the
// symmetrization. Moving the copy by a Cartesian t multiplies F1(k) by exp(2 pi i s_k.t), so the
// derivative of F with respect to a translation of the body comes with it, exactly.
class SymmetryComposition {
public:
    // `grid`: the zone's grid (size, cell and group). `hkl`: the zone's observed indices, sorted. Only
    // those prepare_asu_data(d_min) would give a value on this grid are composed - in the group's
    // reciprocal ASU, strictly inside d_min, inside the grid's Nyquist box, not systematically absent
    // and not 000 - so an observation it does not match stays unmatched, as it was.
    SymmetryComposition(const gemmi::Grid<float> &grid, double d_min, const std::vector<gemmi::Miller> &hkl);

    struct Term {
        size_t index;                 // of F1(hR) in the half-l grid, or of its Friedel mate
        bool conj;                    // read as the Friedel mate: conj F1(-hR)
        std::complex<double> phase;   // exp(+2 pi i h.t)
        gemmi::Vec3 s;                // hR as a Cartesian reciprocal vector
        gemmi::Miller k;              // hR
    };

    // For each of `hkl`, its position among the composed indices, or -1.
    const std::vector<int> &Row() const { return row_; }
    const std::vector<gemmi::Miller> &Hkl() const { return hkl_; }
    // Ops() terms per composed index, index-major, and each index's 1/d^2 and n_cen.
    const std::vector<Term> &Terms() const { return terms_; }
    size_t Ops() const { return ops_; }
    const std::vector<double> &InvD2() const { return inv_d2_; }
    double Centring() const { return centring_; }

    // F at the composed indices from `f1` = MapToFPhi() of the unsymmetrized grid, times the unblur
    // factor prepare_asu_data() applies; with `df_dt`, also dF/dt for a Cartesian translation t.
    void Compose(const gemmi::FPhiGrid<float> &f1, double unblur, std::vector<std::complex<double>> &f,
                 std::vector<std::array<std::complex<double>, 3>> *df_dt, size_t nthreads) const;

private:
    std::vector<int> row_;
    std::vector<gemmi::Miller> hkl_;
    std::vector<double> inv_d2_;
    size_t ops_ = 1;                  // terms per index
    double centring_ = 1;             // n_cen: every centring operator adds the same term
    std::vector<Term> terms_;
};

// What RefineRigidBody's fixed atoms add to a zone: their Fcalc at the zone's composed indices, and the
// bulk-solvent mask of them and the body together (`moving`, where the zone starts) at the same indices.
struct RigidBodyFixedPart {
    std::vector<std::complex<double>> fcalc;
    std::vector<std::complex<float>> fmask;
};
RigidBodyFixedPart RigidBodyFixed(const gemmi::Model &moving, const gemmi::Model &fixed, const gemmi::UnitCell &cell,
                                  const gemmi::SpaceGroup &sg, double d_min, const SymmetryComposition &composition,
                                  size_t nthreads);

// The target RefineRigidBody minimises over one resolution zone: the amplitude residuals of the model
// placed at q (the angle-axis rotation about the centroid times the rms radius, then the translation,
// all in angstroms), with the scale re-fitted at every evaluation. The placement and the counters are
// common; what evaluates the residuals and the Jacobian is the backend: RigidBodyTarget below on the
// CPU, RigidBodyTargetGPU (RigidBodyGPU.h) on a GPU. Both compute the same function.
class RigidBodyTargetBase {
public:
    explicit RigidBodyTargetBase(const gemmi::Model &model);
    virtual ~RigidBodyTargetBase() = default;

    bool Usable() const { return !base_.empty() && rms_radius_ > 0; }
    double RmsRadius() const { return rms_radius_; }
    void Place(const double q[6], gemmi::Model &model) const;

    virtual void SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) = 0;
    virtual size_t NumObservations() const = 0;

    virtual bool Residuals(const double q[6], double *residuals) = 0;
    // d residuals / d q at q, NumObservations() x 6 row-major, with the scale re-fit folded in by
    // projection and the bulk-solvent mask held.
    virtual bool Jacobian(const double q[6], double *jacobian) = 0;

    bool hold_mask = false;        // for the tests: keep the last evaluation's bulk-solvent mask
    const gemmi::Model *fixed = nullptr;   // RefineRigidBody's fixed atoms; set before the first zone
    int evaluations = 0;
    int jacobians = 0;
    int unmatched = 0;             // zone observations with no calculated amplitude to compare against
    double k_sol = 0, b_sol = 0;   // the bulk solvent the zone's target was evaluated with

protected:
    double solvent_b_max_;         // SolventBMax() of the model: the ceiling of the zone's b_sol grid
    std::vector<gemmi::Position> base_;
    gemmi::Position centre_;
    double rms_radius_ = 0;
};

// The CPU backend.
class RigidBodyTarget : public RigidBodyTargetBase {
public:
    // q = 0 is the placement `model` has now. Every evaluation moves `model`.
    RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                    size_t nthreads);

    void SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) override;
    size_t NumObservations() const override { return fobs_.v.size(); }

    bool Residuals(const double q[6], double *residuals) override;
    bool Jacobian(const double q[6], double *jacobian) override;

private:
    void CopyFcalc(const gemmi::Model &model, std::vector<std::complex<double>> &f,
                   std::vector<std::array<std::complex<double>, 3>> *df_dt) const;

    gemmi::Model &model_;
    const gemmi::UnitCell &cell_;
    const gemmi::SpaceGroup &sg_;
    size_t nthreads_;
    std::vector<gemmi::Model> column_models_;   // the rotation columns' own copies

    gemmi::AsuData<gemmi::ValueSigma<float>> fobs_;
    double d_min_ = 0;
    double f_mean_ = 1;
    std::vector<size_t> orbit_leaders_;
    std::optional<SymmetryComposition> composition_;   // the zone's
    RigidBodyFixedPart fixed_part_;                    // the zone's, with fixed atoms
    bool solvent_fitted_ = false;

    // The last evaluation, which the Jacobian is taken at.
    bool have_point_ = false;
    std::array<double, 6> q_{};
    std::vector<std::complex<double>> fc_;
    std::vector<std::array<std::complex<double>, 3>> dfc_dt_;
    std::vector<std::complex<float>> fmask_;
    double k_overall_ = 1;
    gemmi::SMat33<double> b_star_{0, 0, 0, 0, 0, 0};
};