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

#include "RigidBodyRefine.h"

#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <complex>
#include <vector>

#include <ceres/ceres.h>
#include <ceres/rotation.h>

#include "gemmi/dencalc.hpp"   // DensityCalculator
#include "gemmi/it92.hpp"      // IT92 x-ray form factors
#include "gemmi/scaling.hpp"   // Scaling (bulk solvent + anisotropic B)
#include "gemmi/solmask.hpp"   // SolventMasker

#include "ModelFFT.h"               // MapToFPhi
#include "ModelScaling.h"           // FitModelScale
#include "../common/JFJochMath.h"   // PI
#include "../common/Logger.h"
#include "../common/ParallelFor.h"

namespace {

using Table = gemmi::IT92<float>;

// The ladder the placement is walked down. It starts coarse because the model arrives already placed
// but out by a cell's worth of non-isomorphism: at 6 A a few hundred reflections see the body as a
// blob and the target has one broad minimum, and each finer zone starts from the previous one's
// answer. It stops at 3.5 A, which is where rigid-body refinement is conventionally run (it is
// REFMAC's own default through dimple) - the movement being recovered is a few tenths of an
// angstrom, a tenth of that resolution, so it is well determined there, while a finer zone costs
// (1/d)^3 in grid points and reflections for a placement it cannot meaningfully sharpen.
constexpr double LADDER[] = {6.0, 4.5, 3.5};

// The step of the forward-difference Jacobian, as a fraction of the zone's resolution - so it is
// 0.06 A of atom displacement at 6 A and 0.035 A at 3.5 A. A step fixed in angstroms instead is far
// too small for the coarse zones, where a structure factor barely notices it and the derivative is
// swallowed by the jitter of the scale re-fit: measured, a fixed 0.02 A left the 6 A zone at 0.35
// degrees where this rule takes it to 2.79, which is most of the way to the answer.
constexpr double JACOBIAN_STEP_FRACTION = 0.01;

// Parameters are carried as six lengths in angstroms - the first three are the angle-axis rotation
// vector multiplied by the model's rms radius, so a unit of each of the six moves a typical atom by
// the same amount. That makes the Jacobian step isotropic in something physical, rather than mixing
// radians with angstroms.
struct Placement {
    gemmi::Position centre;      // the model centroid: rotating about it decorrelates R from t
    double rms_radius = 1.0;     // rms distance of the atoms from the centroid

    void Apply(const double q[6], const std::vector<gemmi::Position> &base, gemmi::Model &model) const {
        const double aa[3] = {q[0] / rms_radius, q[1] / rms_radius, q[2] / rms_radius};
        size_t i = 0;
        for (gemmi::Chain &ch : model.chains)
            for (gemmi::Residue &r : ch.residues)
                for (gemmi::Atom &a : r.atoms) {
                    const double p[3] = {base[i].x - centre.x, base[i].y - centre.y, base[i].z - centre.z};
                    double rp[3];
                    ceres::AngleAxisRotatePoint(aa, p, rp);
                    a.pos = gemmi::Position(rp[0] + centre.x + q[3],
                                            rp[1] + centre.y + q[4],
                                            rp[2] + centre.z + q[5]);
                    ++i;
                }
    }
};

// One target evaluation: place the model, recompute Fcalc and the bulk-solvent mask to the zone's
// resolution, re-fit the scale, and hand back the amplitude residuals.
class Evaluator {
public:
    Evaluator(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
              const std::vector<gemmi::Position> &base, const Placement &placement, size_t nthreads)
        : model_(model), cell_(cell), sg_(sg), base_(base), placement_(placement), nthreads_(nthreads) {}

    // The zone's observations, and the scale the residuals are expressed in.
    void SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
        fobs_ = fobs;
        d_min_ = d_min;
        double sum = 0;
        for (const auto &hv : fobs_.v)
            sum += hv.value.value;
        f_mean_ = fobs_.v.empty() ? 1.0 : sum / static_cast<double>(fobs_.v.size());
        solvent_fitted_ = false;
    }

    // Hold the zone's bulk solvent at a pair another evaluator already fitted, as if this one had
    // fitted it itself (the Jacobian's column evaluators take the central evaluation's).
    void UseSolvent(double k, double b) {
        k_sol = k;
        b_sol = b;
        solvent_fitted_ = true;
    }

    size_t NumObservations() const { return fobs_.v.size(); }
    double JacobianStep() const { return JACOBIAN_STEP_FRACTION * d_min_; }
    int evaluations = 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

    bool Residuals(const double q[6], double *residuals) {
        ++evaluations;
        placement_.Apply(q, base_, model_);

        gemmi::DensityCalculator<Table, float> dc;
        dc.d_min = d_min_;
        dc.rate = 1.5;
        dc.grid.unit_cell = cell_;
        dc.grid.spacegroup = &sg_;
        dc.set_refmac_compatible_blur(model_);
        dc.put_model_density_on_grid(model_);
        gemmi::AsuData<std::complex<float>> fcalc =
            MapToFPhi(dc.grid).prepare_asu_data(dc.d_min, dc.blur, false, false, false);

        gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
        gemmi::Grid<float> mask_grid;
        mask_grid.unit_cell = cell_;
        mask_grid.spacegroup = &sg_;
        mask_grid.set_size_from_spacing(dc.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
        masker.put_mask_on_grid(mask_grid, model_);
        gemmi::AsuData<std::complex<float>> fmask =
            MapToFPhi(mask_grid).prepare_asu_data(dc.d_min, 0);
        if (fmask.size() != fcalc.size())
            return false;

        // Re-fitted at every evaluation: with the scale held at the starting placement's value the
        // target would measure the scale as much as the placement, and the body would translate to
        // repair a scale error instead of moving where the density is.
        //
        // The bulk solvent is not part of that scale. k_sol and b_sol describe the disordered solvent
        // of the crystal rather than the fit of one placement, so they are fitted once per zone - by
        // FitModelScale, inside the same physical box the reported fit is searched in - and then held
        // while the overall scale and the anisotropic B follow the body. Leaving them free at every
        // evaluation, which is what gemmi's unbounded Levenberg-Marquardt did here, puts a solvent
        // term with no physical meaning inside the target that decides where the model goes: measured
        // over a corpus of deposited models, 40% of the evaluations came out with b_sol outside
        // 10-80 A^2, some of them negative, which is a solvent that GROWS with resolution.
        gemmi::Scaling<float> scaling(cell_, &sg_);
        scaling.use_solvent = true;
        scaling.prepare_points(fcalc, fobs_, &fmask);
        if (scaling.points.empty())
            return false;
        if (!solvent_fitted_) {
            FitModelScale(scaling, {}, nthreads_);
            k_sol = scaling.k_sol;
            b_sol = scaling.b_sol;
            solvent_fitted_ = true;
        }
        scaling.k_sol = k_sol;
        scaling.b_sol = b_sol;
        scaling.fix_k_sol = true;
        scaling.fix_b_sol = true;
        scaling.fit_isotropic_b_approximately();
        scaling.fit_parameters();
        scaling.scale_data(fcalc, &fmask);

        // Both are sorted and in the same ASU, so one merge pass matches them. An observation with no
        // calculated amplitude gets residual 0, which drops it from the target rather than scoring it
        // as a perfect fit: its Jacobian row below comes out zero as well, and the set that matches is
        // fixed by the cell, the group and the zone, so it does not move as the body does. It is
        // counted and reported because a large count is a statement about the model rather than about
        // this refinement - a group whose reflection conditions the data do not obey leaves half of
        // them with nothing to compare against.
        auto c = fcalc.v.begin();
        unmatched = 0;
        for (size_t i = 0; i < fobs_.v.size(); ++i) {
            const gemmi::Miller &h = fobs_.v[i].hkl;
            while (c != fcalc.v.end() && c->hkl < h)
                ++c;
            const bool matched = c != fcalc.v.end() && c->hkl == h;
            if (!matched)
                ++unmatched;
            residuals[i] = matched ? (fobs_.v[i].value.value - std::abs(c->value)) / f_mean_ : 0.0;
        }
        return true;
    }

private:
    gemmi::Model &model_;
    const gemmi::UnitCell &cell_;
    const gemmi::SpaceGroup &sg_;
    const std::vector<gemmi::Position> &base_;
    Placement placement_;
    gemmi::AsuData<gemmi::ValueSigma<float>> fobs_;
    double d_min_ = 0;
    double f_mean_ = 1;
    bool solvent_fitted_ = false;
    size_t nthreads_ = 1;
};

// Ceres' own numeric differentiation steps by |x| * relative_step_size, which is zero at the start of
// every zone (the placement begins at no shift), so the Jacobian is supplied here instead, by
// forward differences at a step chosen in the parameters' units. Analytic dF/dp would need
// derivatives GEMMI's structure-factor path does not have, and at six parameters it is not worth it:
// a Jacobian costs seven evaluations, and the evaluations at 6-3.5 A are cheap.
//
// The six shifted evaluations are independent of each other, so they run in parallel, column j on
// `columns[j]`: an evaluator over its own copy of the model (an evaluation moves every atom to the
// placement it is asked about, so two cannot share one), set to the same zone and handed the solvent
// pair the central evaluation fitted - which is the pair the serial loop's shifted evaluations used,
// since the zone's solvent is fitted by the first evaluation of the zone and every Evaluate starts
// with the central one. Each column's arithmetic is the serial one, so the Jacobian is too.
class RigidBodyCost : public ceres::CostFunction {
public:
    RigidBodyCost(Evaluator &ev, std::vector<Evaluator> &columns, size_t nthreads)
        : ev_(ev), columns_(columns), nthreads_(nthreads) {
        set_num_residuals(static_cast<int>(ev.NumObservations()));
        mutable_parameter_block_sizes()->push_back(6);
    }

    bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const override {
        const int n = num_residuals();
        if (!ev_.Residuals(parameters[0], residuals))
            return false;
        if (jacobians != nullptr && jacobians[0] != nullptr) {
            const double step = ev_.JacobianStep();
            std::vector<std::vector<double>> shifted(6, std::vector<double>(n));
            std::array<bool, 6> ok{};
            ParallelFor(6, nthreads_, [&](int j) {
                double q[6];
                std::copy(parameters[0], parameters[0] + 6, q);
                q[j] += step;
                columns_[j].UseSolvent(ev_.k_sol, ev_.b_sol);
                ok[j] = columns_[j].Residuals(q, shifted[j].data());
            });
            // Counted as the serial loop counted them: it stopped at the first column that failed.
            for (int j = 0; j < 6; j++) {
                ++ev_.evaluations;
                ev_.unmatched = columns_[j].unmatched;
                if (!ok[j])
                    return false;
                for (int i = 0; i < n; i++)
                    jacobians[0][i * 6 + j] = (shifted[j][i] - residuals[i]) / step;
            }
        }
        return true;
    }

private:
    Evaluator &ev_;
    std::vector<Evaluator> &columns_;
    size_t nthreads_;
};

// The directions in which this space group's origin is free. Translating the whole cell content along
// one of them multiplies every F by a phase and leaves every |F| EXACTLY unchanged, so the target
// cannot determine that component: all three directions in P1, the unique axis in a polar group. The
// R-free gate cannot stand in for this - it is a function of |F| too, so along such a direction it
// sees only grid noise and commits or not by coin flip, while the other five parameters carry the
// noise in with them. The free directions are the common fixed subspace of the group's rotation
// parts, and the projector onto it is simply their average.
gemmi::Mat33 GaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell) {
    const gemmi::GroupOps gops = sg.operations();
    double m[3][3] = {};
    const double n = static_cast<double>(gops.sym_ops.size()) * gemmi::Op::DEN;
    for (const gemmi::Op &op : gops.sym_ops)
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                m[i][j] += static_cast<double>(op.rot[i][j]) / n;
    const gemmi::Mat33 mean(m[0][0], m[0][1], m[0][2],
                            m[1][0], m[1][1], m[1][2],
                            m[2][0], m[2][1], m[2][2]);
    // Fractional projector taken into orthogonal space, where the parameters live.
    return cell.orth.mat.multiply(mean).multiply(cell.frac.mat);
}

}  // namespace

std::vector<gemmi::Position> ModelPositions(const gemmi::Model &model) {
    std::vector<gemmi::Position> pos;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                pos.push_back(a.pos);
    return pos;
}

void SetModelPositions(gemmi::Model &model, const std::vector<gemmi::Position> &pos) {
    size_t i = 0;
    for (gemmi::Chain &ch : model.chains)
        for (gemmi::Residue &r : ch.residues)
            for (gemmi::Atom &a : r.atoms)
                a.pos = pos[i++];
}

// One rigid body, not groups: a fragment-screening model arrives already solved and isomorphous, and
// the movement to recover is the crystal's, not the molecule's. Splitting it into domains or giving a
// bound ligand its own six parameters would refine against evidence this data does not separately
// carry, and the ligand is what the difference map is meant to show rather than model away.
//
// Some of the translation would be a gauge rather than a quantity - the origin is free in all three
// directions in P1 and along the unique axis in a polar group, and |F| does not change when the whole
// content moves along it - so that component is projected out after every zone. Neither of the two
// things that might look like they cover it actually does: the R-free gate is a function of |F| and
// therefore blind to exactly this, and the LM damping follows the gauge column of the Jacobian, which
// is not zero but noise divided by the difference step.
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) {
    const auto t0 = std::chrono::steady_clock::now();
    RigidBodyRefineResult result;

    const std::vector<gemmi::Position> base = ModelPositions(model);
    if (base.empty() || fobs.v.empty())
        return result;

    Placement placement;
    for (const gemmi::Position &p : base)
        placement.centre += p;
    placement.centre *= 1.0 / static_cast<double>(base.size());
    double r2 = 0;
    for (const gemmi::Position &p : base)
        r2 += placement.centre.dist_sq(p);
    placement.rms_radius = std::sqrt(r2 / static_cast<double>(base.size()));
    if (!(placement.rms_radius > 0))
        return result;

    std::vector<double> ladder;
    for (double zone : LADDER)
        if (zone >= d_min)
            ladder.push_back(zone);
    if (ladder.empty())
        ladder.push_back(d_min);

    const gemmi::Mat33 gauge = GaugeProjector(sg, cell);
    Evaluator ev(model, cell, sg, base, placement, nthreads);
    // The Jacobian's six column evaluators, each over its own copy of the model (see RigidBodyCost).
    // The copies only ever hold probe placements; `model` is left at the answer below.
    std::vector<gemmi::Model> column_models(6, model);
    std::vector<Evaluator> columns;
    columns.reserve(6);
    for (int j = 0; j < 6; j++)
        columns.emplace_back(column_models[j], cell, sg, base, placement, 1);
    double q[6] = {0, 0, 0, 0, 0, 0};
    bool any_zone_solved = false;
    for (double zone : ladder) {
        gemmi::AsuData<gemmi::ValueSigma<float>> zone_obs;
        zone_obs.unit_cell_ = fobs.unit_cell_;
        zone_obs.spacegroup_ = fobs.spacegroup_;
        for (const auto &hv : fobs.v)
            if (cell.calculate_d(hv.hkl) >= zone)
                zone_obs.v.push_back(hv);
        if (zone_obs.v.size() < 50)
            continue;
        result.zones.push_back(zone);   // the ladder WALKED, which a thin zone drops out of
        ev.SetZone(zone_obs, zone);
        for (Evaluator &c : columns)
            c.SetZone(zone_obs, zone);

        ceres::Problem problem;
        problem.AddResidualBlock(new RigidBodyCost(ev, columns, nthreads), nullptr, q);
        ceres::Solver::Options options;
        options.linear_solver_type = ceres::DENSE_QR;
        options.max_num_iterations = 15;
        options.function_tolerance = 1e-4;
        options.parameter_tolerance = 1e-4;
        options.logging_type = ceres::LoggingType::SILENT;
        ceres::Solver::Summary summary;
        const int evaluations_before = ev.evaluations;
        const auto zone_t0 = std::chrono::steady_clock::now();
        ceres::Solve(options, &problem, &summary);
        any_zone_solved = any_zone_solved || summary.IsSolutionUsable();
        const gemmi::Vec3 along = gauge.multiply(gemmi::Vec3(q[3], q[4], q[5]));
        q[3] -= along.x; q[4] -= along.y; q[5] -= along.z;
        logger.Debug("Rigid body zone {:.1f} A: {} reflections ({} without a calculated amplitude), "
                     "solvent k_sol {:.2f} b_sol {:.0f} A^2, {} iterations, {} evaluations, {:.2f} s, "
                     "rotation {:.3f} deg, translation {:.3f} A", zone, zone_obs.v.size(), ev.unmatched,
                     ev.k_sol, ev.b_sol,
                     summary.iterations.empty() ? 0 : summary.iterations.size() - 1,
                     ev.evaluations - evaluations_before,
                     std::chrono::duration<double>(std::chrono::steady_clock::now() - zone_t0).count(),
                     std::sqrt(q[0]*q[0] + q[1]*q[1] + q[2]*q[2]) / placement.rms_radius * 180.0 / PI,
                     std::sqrt(q[3]*q[3] + q[4]*q[4] + q[5]*q[5]));
    }

    placement.Apply(q, base, model);   // Ceres left the model at a Jacobian probe; put it at the answer
    result.evaluations = ev.evaluations;
    result.converged = any_zone_solved;
    result.k_sol = ev.k_sol;
    result.b_sol = ev.b_sol;
    const double aa = std::sqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2]) / placement.rms_radius;
    result.angle_deg = aa * 180.0 / PI;
    result.shift_A = std::sqrt(q[3] * q[3] + q[4] * q[4] + q[5] * q[5]);
    result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();

    if (!result.zones.empty())
        logger.Info("Model validation: rigid body over {} resolution zone(s) down to {:.1f} A, "
                    "{} evaluations in {:.2f} s: rotation {:.3f} deg, translation {:.3f} A",
                    result.zones.size(), result.zones.back(), result.evaluations, result.seconds,
                    result.angle_deg, result.shift_A);
    else
        logger.Info("Model validation: rigid body had no resolution zone with enough reflections to "
                    "run in; the model is left where it arrived");
    return result;
}