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

// FixedSolventFit follows the scaling target of GEMMI's scaling.hpp
// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/scaling.hpp), and
// ComputeLMMatrices / LevMarFit the Levenberg-Marquardt of GEMMI's levmar.hpp
// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/levmar.hpp)
// (c) Global Phasing Ltd., Mozilla Public License Version 2.0

#include "ModelScaling.h"

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

#include "../../common/ParallelFor.h"

namespace {

// gemmi::Scaling<float> as a grid point fits it: the same starting point, parameters, model values,
// derivatives and R. The solvent pair is fixed there, so |Fcalc + k_sol exp(-b_sol s^2) Fmask| of a
// reflection is the same at every evaluation of the solver; it is taken once per grid point here rather than
// at every evaluation, where it was most of the cost of the fit. Every expression is the one
// scaling.hpp evaluates, in the same types, so every number is the same to the bit.
struct FixedSolventFit {
    struct Point {
        gemmi::Miller hkl;
        float fobs;
        float fcalc_abs;   // std::abs(Scaling::get_fcalc(p))
        double get_y() const { return fobs; }
        double get_weight() const { return 1.0; }
    };
    const std::vector<gemmi::Vec6> &constraint_matrix;
    double k_overall = 1.0;
    gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
    std::vector<Point> points;

    explicit FixedSolventFit(const gemmi::Scaling<float> &s) : constraint_matrix(s.constraint_matrix) {}

    // The reflections of `s` at the solvent pair (k_sol, b_sol), and the start of the fit there: what
    // Scaling::fit_isotropic_b_approximately() sets, from the same |Fcalc|, taken once for both. Where
    // that finds five or fewer reflections to fit on it sets nothing, and the fit starts from
    // (k_start, b_start) - the scale gemmi's would have been left at.
    void Start(const gemmi::Scaling<float> &s, double k_sol, double b_sol,
               double k_start, const gemmi::SMat33<double> &b_start) {
        k_overall = k_start;
        b_star = b_start;
        points.resize(s.points.size());
        double sx = 0, sy = 0, sxx = 0, sxy = 0;
        int n = 0;
        for (size_t i = 0; i < points.size(); ++i) {
            const auto &p = s.points[i];
            // Scaling::get_fcalc() at (k_sol, b_sol)
            const std::complex<float> fcalc =
                s.use_solvent ? p.fcmol + (float) (k_sol * std::exp(-b_sol * p.stol2)) * p.fmask : p.fcmol;
            points[i] = {p.hkl, p.fobs, std::abs(fcalc)};
            if (p.fobs < 1 || p.fobs < p.sigma)  // skip weak reflections
                continue;
            double fcalc_abs = points[i].fcalc_abs;
            double x = p.stol2;
            double y = std::log(static_cast<float>(p.fobs / fcalc_abs));
            sx += x;
            sy += y;
            sxx += x * x;
            sxy += x * y;
            n += 1;
        }
        if (n <= 5)
            return;
        double slope = (n * sxy - sx * sy) / (n * sxx - sx * sx);
        double intercept = (sy - slope * sx) / n;
        double b_iso = -slope;
        k_overall = std::exp(intercept);
        b_star = gemmi::SMat33<double>{b_iso, b_iso, b_iso, 0, 0, 0}.transformed_by(s.cell.frac.mat);
    }

    // Scaling::get_parameters(), set_parameters(), get_overall_scale_factor(), compute_value() and
    // compute_value_and_derivatives(), with k_sol and b_sol fixed.
    std::vector<double> get_parameters() const {
        std::vector<double> ret;
        ret.push_back(k_overall);
        for (const gemmi::Vec6 &v : constraint_matrix)
            ret.push_back(gemmi::vec6_dot(v, b_star));
        return ret;
    }
    void set_parameters(const double *p) {
        k_overall = p[0];
        int n = 0;
        b_star = {0, 0, 0, 0, 0, 0};
        for (const gemmi::Vec6 &row : constraint_matrix) {
            double d = p[++n];
            b_star.u11 += row[0] * d;
            b_star.u22 += row[1] * d;
            b_star.u33 += row[2] * d;
            b_star.u12 += row[3] * d;
            b_star.u13 += row[4] * d;
            b_star.u23 += row[5] * d;
        }
    }
    void set_parameters(const std::vector<double> &p) { set_parameters(p.data()); }
    double get_overall_scale_factor(const gemmi::Miller &hkl) const {
        return k_overall * std::exp(-0.25 * b_star.r_u_r(hkl));
    }
    double compute_value(const Point &p) const {
        return p.fcalc_abs * (float) get_overall_scale_factor(p.hkl);
    }
    // For NA parameters: the B components are NA - 1 constraint rows.
    template <int NA>
    double compute_value_and_derivatives(const Point &p, double *dy_da) const {
        gemmi::Vec3 h(p.hkl);
        double kaniso = std::exp(-0.25 * b_star.r_u_r(h));
        double fcalc_abs = p.fcalc_abs;
        int n = 1;
        double fe = fcalc_abs * kaniso;
        double y = k_overall * fe;
        dy_da[0] = fe;
        gemmi::SMat33<double> du = {
            -0.25 * y * (h.x * h.x),
            -0.25 * y * (h.y * h.y),
            -0.25 * y * (h.z * h.z),
            -0.5 * y * (h.x * h.y),
            -0.5 * y * (h.x * h.z),
            -0.5 * y * (h.y * h.z),
        };
        for (int j = 0; j < NA - 1; ++j)
            dy_da[n + j] = gemmi::vec6_dot(constraint_matrix[j], du);
        return y;
    }
};

// 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 FixedSolventFit &fit) {
    double num = 0, den = 0;
    for (const auto &p : fit.points) {
        num += std::fabs(p.fobs - fit.compute_value(p));
        den += p.fobs;
    }
    return den > 0 ? num / den : 1.0;
}

// gemmi's compute_lm_matrices() (levmar.hpp) on a FixedSolventFit of NA parameters. With the count
// known to the compiler the matrices are held in registers rather than written to memory at every
// reflection; the operations are gemmi's, in gemmi's order, so the matrices are gemmi's to the bit.
// A reflection's weight is 1 (FixedSolventFit::Point::get_weight), which gemmi multiplies by.
// gemmi skips the row of a derivative that is exactly 0; here the row gets +0 instead, which keeps the
// loop free of branches and changes nothing: an element of alpha or beta is never -0, the one value
// adding +0 would change, since it starts at +0 and a sum is only -0 when both its terms are. And
// where gemmi fills the lower half of alpha, the whole square is summed here, a rectangle the compiler
// can vectorise; the lower half is gemmi's, and the upper is replaced by it at the end, as in gemmi.
template <int NA>
double ComputeLMMatrices(const FixedSolventFit &fit, double *alpha_out, double *beta_out) {
    long double wssr = 0; // long double here notably increases the accuracy
    double alpha[NA * NA] = {}, beta[NA] = {};
    double dy_da[NA];
    for (const auto &p : fit.points) {
        double y = fit.compute_value_and_derivatives<NA>(p, dy_da);
        double dy_sig = p.get_y() - y;
        for (int j = 0; j != NA; ++j) {
            const bool used = dy_da[j] != 0;
            for (int k = 0; k < NA; ++k) {
                const double a = dy_da[j] * dy_da[k];
                alpha[NA * j + k] += used ? a : 0.0;
            }
            const double b = dy_sig * dy_da[j];
            beta[j] += used ? b : 0.0;
        }
        wssr += gemmi::sq(dy_sig);
    }
    // The upper half of alpha is the lower half's mirror, as in gemmi.
    for (int j = 1; j < NA; j++)
        for (int k = 0; k < j; k++)
            alpha[NA * k + j] = alpha[NA * j + k];
    std::copy(alpha, alpha + NA * NA, alpha_out);
    std::copy(beta, beta + NA, beta_out);
    return (double) wssr;
}

// gemmi's LevMar::fit() (levmar.hpp) with its default settings, on ComputeLMMatrices<NA>() - the same
// iterations to the same answer, bit for bit.
template <int NA>
void LevMarFit(FixedSolventFit &target) {
    const int eval_limit = 100;
    const double lambda_limit = 1e+15;
    const double stop_rel_change = 1e-5;
    const double lambda_up_factor = 10;
    const double lambda_down_factor = 0.1;
    const double lambda_start = 0.001;

    std::vector<double> initial_a = target.get_parameters();
    std::vector<double> best_a = initial_a;
    double lambda = lambda_start;
    double alpha[NA * NA], beta[NA], temp_alpha[NA * NA];
    std::vector<double> temp_beta(NA);

    const double initial_wssr = ComputeLMMatrices<NA>(target, alpha, beta);
    double wssr = initial_wssr;

    int small_change_counter = 0;
    int eval_count = 1;  // number of function evaluations so far
    for (;;) {
        if (eval_limit > 0 && eval_count >= eval_limit)
            break;

        // prepare next parameters -> temp_beta
        std::copy(alpha, alpha + NA * NA, temp_alpha);
        // Using '*=' not '+=' below applies the dampling factor as:
        // J^T J + lambda * diag(J^T J); not ... + lambda * I.
        for (int j = 0; j < NA; j++)
            temp_alpha[NA * j + j] *= (1.0 + lambda);
        std::copy(beta, beta + NA, temp_beta.begin());

        // Matrix solution (Ax=b)  temp_alpha * da == temp_beta
        gemmi::jordan_solve(temp_alpha, temp_beta.data(), NA);

        for (int i = 0; i < NA; i++)
            // put new a[] into temp_beta[]
            temp_beta[i] += best_a[i];

        target.set_parameters(temp_beta);
        double new_wssr = gemmi::compute_wssr(target);
        ++eval_count;
        if (new_wssr < wssr) {
            double rel_change = (wssr - new_wssr) / wssr;
            wssr = new_wssr;
            best_a = temp_beta;

            if (wssr == 0)
                break;
            // termination criterion: negligible change of wssr
            if (rel_change < stop_rel_change) {
                if (++small_change_counter >= 2)
                    break;
            } else {
                small_change_counter = 0;
            }
            ComputeLMMatrices<NA>(target, alpha, beta);
            ++eval_count;
            lambda *= lambda_down_factor;
        } else { // worse fitting
            if (lambda > lambda_limit) // termination criterion: large lambda
                break;
            lambda *= lambda_up_factor;
        }
    }

    target.set_parameters(wssr < initial_wssr ? best_a : initial_a);
}

// LevMarFit() at the fit's parameter count: k_overall and one per symmetry-allowed component of B,
// which is one (cubic) to six (triclinic).
void LevMarFit(FixedSolventFit &fit) {
    switch (fit.constraint_matrix.size()) {
        case 1: LevMarFit<2>(fit); break;
        case 2: LevMarFit<3>(fit); break;
        case 3: LevMarFit<4>(fit); break;
        case 4: LevMarFit<5>(fit); break;
        case 5: LevMarFit<6>(fit); break;
        default: LevMarFit<7>(fit); break;
    }
}

}  // 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.
//
// Every grid point is its own fit: fit_isotropic_b_approximately() sets k_overall and b_star from the
// data and the point's solvent pair alone, so a point does not depend on the one fitted before it, and
// the points run in parallel, one point per task, each on a FixedSolventFit of its own. The winner is
// then read off in grid order with the serial rule (lowest finite R, the first on a tie), so the answer
// is the serial loop's bit for bit. For the same reason a pair the refinement pass shares with the
// coarse one - the same two numbers, to the bit - is the same fit, and is taken from there rather than
// made again. The one exception is fit_isotropic_b_approximately() finding five or fewer reflections to
// fit on - it then returns without setting anything and a point WOULD start from where the previous one
// ended - so there the grid is walked in order, each point from where the last ended, as it always was.
ModelScaleReport FitModelScale(gemmi::Scaling<float> &scaling, ModelScaleBox box, size_t nthreads) {
    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;

    // The reflections fit_isotropic_b_approximately() fits on (its own filter).
    int n_isotropic = 0;
    for (const auto &p : scaling.points)
        if (!(p.fobs < 1 || p.fobs < p.sigma))
            ++n_isotropic;
    const bool independent = n_isotropic > 5;

    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};

    struct PointFit { double k_sol, b_sol, r = NAN, k_overall = 1.0; gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0}; };
    // The scale the walk stands at, which a point starts from only where the data give it no start of its own.
    double k_last = scaling.k_overall;
    gemmi::SMat33<double> b_last = scaling.b_star;
    auto fit_point = [&scaling](FixedSolventFit &fit, PointFit &pf, double k_start,
                                const gemmi::SMat33<double> &b_start) {
        fit.Start(scaling, pf.k_sol, pf.b_sol, k_start, b_start);
        LevMarFit(fit);                      // k_overall + anisotropic B only
        pf.r = RFactor(fit);
        pf.k_overall = fit.k_overall;
        pf.b_star = fit.b_star;
    };
    std::vector<PointFit> fitted;            // every point fitted so far, in grid order
    auto try_points = [&](std::vector<PointFit> &pts) {
        if (independent) {
            std::vector<int> todo;
            for (int i = 0; i < static_cast<int>(pts.size()); ++i) {
                const auto same = std::find_if(fitted.begin(), fitted.end(), [&](const PointFit &f) {
                    return f.k_sol == pts[i].k_sol && f.b_sol == pts[i].b_sol;
                });
                if (same != fitted.end())
                    pts[i] = *same;
                else
                    todo.push_back(i);
            }
            ParallelFor(static_cast<int>(todo.size()), nthreads, [&](int j) {
                FixedSolventFit fit(scaling);
                fit_point(fit, pts[todo[j]], k_last, b_last);
            });
        } else {
            FixedSolventFit fit(scaling);
            for (auto &pf : pts) {
                fit_point(fit, pf, k_last, b_last);
                k_last = pf.k_overall;
                b_last = pf.b_star;
            }
        }
        for (const auto &pf : pts) {
            ++report.n_grid;
            // A diverged fit gives r = NaN; latched as best_r it wins every later r < best_r.
            if (std::isfinite(pf.r) && (best_r < 0 || pf.r < best_r)) {
                best_r = pf.r;
                best_k_sol = pf.k_sol;
                best_b_sol = pf.b_sol;
                best_k_overall = pf.k_overall;
                best_b_star = pf.b_star;
            }
        }
        fitted.insert(fitted.end(), pts.begin(), pts.end());
    };

    // Coarse pass over the whole box, then one refinement pass around the winner.
    std::vector<PointFit> coarse;
    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)
            coarse.push_back(PointFit{ks, bs});
    try_points(coarse);
    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);
    std::vector<PointFit> fine;
    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)
            fine.push_back(PointFit{ks, bs});
    try_points(fine);

    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;
}