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

#include <catch2/catch_all.hpp>

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

#include "../rugnux/ModelScaling.h"
#include "gemmi/scaling.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

namespace {

// A synthetic set of scaling points: |Fobs| is generated from a known scale, isotropic B, k_sol
// and b_sol, so the fit has a right answer to find. The "structure factors" are a hash of the
// index, so the set is reproduced bit for bit.
void MakePoints(gemmi::Scaling<float> &scaling, const gemmi::UnitCell &cell,
                double k_overall, double b_iso, double k_sol, double b_sol,
                double d_min, double d_max) {
    scaling.points.clear();
    const int hmax = static_cast<int>(cell.a / d_min) + 1;
    const int kmax = static_cast<int>(cell.b / d_min) + 1;
    const int lmax = static_cast<int>(cell.c / d_min) + 1;
    uint32_t seed = 12345;
    auto rnd = [&seed]() {
        seed = seed * 1664525u + 1013904223u;
        return static_cast<double>((seed >> 8) & 0xffff) / 65535.0;
    };
    for (int h = 0; h <= hmax; ++h)
        for (int k = 0; k <= kmax; ++k)
            for (int l = 0; l <= lmax; ++l) {
                if (h == 0 && k == 0 && l == 0) continue;
                const gemmi::Miller hkl{{h, k, l}};
                const double d = cell.calculate_d(hkl);
                if (d < d_min || d > d_max) continue;
                const double stol2 = cell.calculate_stol_sq(hkl);
                const std::complex<double> fc(20.0 + 80.0 * rnd(), 0.0);
                const std::complex<double> fm(5.0 + 5.0 * rnd(), 0.0);
                const std::complex<double> total = fc + k_sol * std::exp(-b_sol * stol2) * fm;
                const double fobs = k_overall * std::exp(-b_iso * stol2) * std::abs(total);
                gemmi::Scaling<float>::Point p{};
                p.hkl = hkl;
                p.stol2 = stol2;
                p.fcmol = std::complex<float>(static_cast<float>(fc.real()), 0.f);
                p.fmask = std::complex<float>(static_cast<float>(fm.real()), 0.f);
                p.fobs = static_cast<float>(fobs);
                p.sigma = 1.f;
                scaling.points.push_back(p);
            }
}

gemmi::UnitCell Cell() {
    gemmi::UnitCell c;
    c.set(60.0, 70.0, 80.0, 90.0, 90.0, 90.0);
    return c;
}

}  // namespace

// The fit has to find a solvent pair it was given, not merely land somewhere legal.
TEST_CASE("ModelScaling recovers a known bulk solvent") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/2.0, /*d_max=*/50.0);
    REQUIRE(scaling.points.size() > 500);

    const ModelScaleReport report = FitModelScale(scaling);
    CHECK(report.n_grid > 0);
    // The grid steps are 0.025 in k_sol and 5 in b_sol, so this is one step of tolerance.
    CHECK(scaling.k_sol == Catch::Approx(0.35).margin(0.03));
    CHECK(scaling.b_sol == Catch::Approx(45.0).margin(6.0));
    CHECK(report.r_work_fit < 0.05);
}

// The regression this exists for: gemmi's own fit_parameters() is an unbounded Levenberg-Marquardt
// and, where the data cannot determine the solvent, walks k_sol/b_sol out of the range a flat
// solvent model means anything in (b_sol of 1707 A^2 was measured on real data). Here the solvent
// contribution is made unidentifiable - only high-resolution data, where exp(-b_sol * s^2) is
// indistinguishable from zero for any large b_sol - and the fit must still return a physical pair.
TEST_CASE("ModelScaling stays physical when the solvent is unidentifiable") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    // Nothing below 3 A: the bulk solvent is a low-resolution feature, so it has almost
    // no leverage here.
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/1.2, /*d_max=*/3.0);
    REQUIRE(scaling.points.size() > 500);

    FitModelScale(scaling);
    CHECK(scaling.k_sol >= 0.15);
    CHECK(scaling.k_sol <= 0.50);
    CHECK(scaling.b_sol >= 10.0);
    CHECK(scaling.b_sol <= 80.0);

    // ... where gemmi's own unbounded fit is free to leave that range.
    gemmi::Scaling<float> unbounded(cell, sg);
    unbounded.use_solvent = true;
    unbounded.points = scaling.points;
    unbounded.fit_isotropic_b_approximately();
    unbounded.fit_parameters();
    CHECK(std::isfinite(unbounded.b_sol));   // it converges; it is simply not bounded
}

// A densely packed crystal with no disordered solvent channels - a small molecule, say - has an
// empty solvent mask, and then Fmask is zero for every reflection. The bulk-solvent contribution
// k_sol * exp(-b_sol s^2) * Fmask is then identically zero WHATEVER k_sol and b_sol come out as,
// so a solvent term fitted where there is no solvent cannot add anything to the model. That is why
// this needs no solvent-content threshold and no switch: the mask already carries the answer.
TEST_CASE("ModelScaling adds nothing when there is no solvent to model") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/2.0, /*d_max=*/50.0);
    REQUIRE(scaling.points.size() > 500);
    for (auto &p : scaling.points)     // an empty mask: no solvent-accessible volume
        p.fmask = {0.f, 0.f};

    FitModelScale(scaling);
    // Whatever the search settled on, the model it produces is the solvent-free one.
    double worst = 0;
    for (const auto &p : scaling.points)
        worst = std::max(worst, static_cast<double>(
                            std::fabs(std::abs(scaling.get_fcalc(p)) - std::abs(p.fcmol))));
    CHECK(worst == Catch::Approx(0.0).margin(1e-6));
    // and the parameters are still reported inside the physical box, not at some arbitrary value
    CHECK(scaling.k_sol >= 0.10);
    CHECK(scaling.k_sol <= 0.60);
    CHECK(scaling.b_sol >= 10.0);
    CHECK(scaling.b_sol <= 80.0);
}