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

#include "ModelScaling.h"

#include <algorithm>
#include <cmath>

namespace {

// R-factor of the current parameters over the fitted reflections. This is what the grid is
// selected on, and it is the quantity the scale exists to make small.
double RFactor(const gemmi::Scaling<float> &scaling) {
    double num = 0, den = 0;
    for (const auto &p : scaling.points) {
        num += std::fabs(p.fobs - scaling.compute_value(p));
        den += p.fobs;
    }
    return den > 0 ? num / den : 1.0;
}

}  // namespace

// Following the phenix bulk-solvent and scaling procedure: k_sol and b_sol by a grid search, with
// the overall scale and the anisotropic B refitted at every grid point - Afonine, Grosse-Kunstleve
// & Adams, Acta Cryst. D61, 850-855, 2005, which searches b_sol over 10-80 A^2 in steps of 5.
// The fit is unweighted, as in both phenix and Refmac (Murshudov, Skubak, Lebedev, Pannu, Steiner,
// Nicholls, Winn, Long & Vagin, Acta Cryst. D67, 355-367, 2011, eq. 11). The physical range and the
// starting values are those of Fokine & Urzhumtsev, Acta Cryst. D58, 1387-1392, 2002.
//
// The point of the grid is that k_sol and b_sol cannot leave the physical box: gemmi's own
// fit_parameters() is an unbounded Levenberg-Marquardt, and on this corpus it reached b_sol of
// 1707 A^2 - a solvent term switched off in all but the lowest-resolution shell. Here the solvent
// pair is held fixed at each grid point and only the overall scale and the symmetry-constrained
// anisotropic B are refined, which is the well-conditioned half of the problem and is left to
// gemmi's solver rather than reimplemented.
ModelScaleReport FitModelScale(gemmi::Scaling<float> &scaling, ModelScaleBox box) {
    ModelScaleReport report;
    report.n_points = static_cast<int>(scaling.points.size());
    if (scaling.points.size() < 20)
        return report;

    const bool had_solvent = scaling.use_solvent;
    scaling.fix_k_sol = true;      // the grid owns the solvent pair; the solver never sees it
    scaling.fix_b_sol = true;

    double best_r = -1, best_k_sol = 0.35, best_b_sol = 46.0, best_k_overall = 1.0;
    gemmi::SMat33<double> best_b_star{0, 0, 0, 0, 0, 0};

    auto try_point = [&](double k_sol, double b_sol) {
        scaling.k_sol = k_sol;
        scaling.b_sol = b_sol;
        scaling.fit_isotropic_b_approximately();   // a fresh starting point for this solvent pair
        scaling.fit_parameters();                  // k_overall + anisotropic B only
        ++report.n_grid;
        const double r = RFactor(scaling);
        // A diverged fit gives r = NaN; latched as best_r it wins every later r < best_r.
        if (std::isfinite(r) && (best_r < 0 || r < best_r)) {
            best_r = r;
            best_k_sol = k_sol;
            best_b_sol = b_sol;
            best_k_overall = scaling.k_overall;
            best_b_star = scaling.b_star;
        }
    };

    // Coarse pass over the whole box, then one refinement pass around the winner.
    for (double ks = box.k_lo; ks <= box.k_hi + 1e-9; ks += 0.05)
        for (double bs = box.b_lo; bs <= box.b_hi + 1e-9; bs += 10.0)
            try_point(ks, bs);
    const double k0 = best_k_sol, b0 = best_b_sol;
    const double k_hi2 = std::min(box.k_hi, k0 + 0.05);
    const double b_hi2 = std::min(box.b_hi, b0 + 10.0);
    for (double ks = std::max(box.k_lo, k0 - 0.05); ks <= k_hi2 + 1e-9; ks += 0.025)
        for (double bs = std::max(box.b_lo, b0 - 10.0); bs <= b_hi2 + 1e-9; bs += 5.0)
            try_point(ks, bs);

    scaling.k_sol = best_k_sol;
    scaling.b_sol = best_b_sol;
    scaling.k_overall = best_k_overall;
    scaling.b_star = best_b_star;
    scaling.use_solvent = had_solvent;
    report.r_work_fit = best_r;
    return report;
}