// 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 <cstdio>
#include <filesystem>
#include <fstream>
#include <random>

#include <gemmi/mmread_gz.hpp>

#include "../common/Logger.h"
#include "../rugnux/ModelValidation.h"
#include "../rugnux/RigidBodyRefine.h"
#include "../rugnux/SigmaA.h"

namespace {
    // A synthetic P1 cell with two carbon atoms - enough for a reader to produce a Structure with
    // atoms, a cell and a space group, which is all these tests are about. Neutral by construction:
    // no real specimen's parameters are involved.
    const char *kMmcif = R"(data_test
_cell.length_a      40.000
_cell.length_b      50.000
_cell.length_c      60.000
_cell.angle_alpha   90.000
_cell.angle_beta    90.000
_cell.angle_gamma   90.000
_symmetry.space_group_name_H-M   'P 1'
loop_
_atom_site.group_PDB
_atom_site.id
_atom_site.type_symbol
_atom_site.label_alt_id
_atom_site.label_atom_id
_atom_site.label_comp_id
_atom_site.label_asym_id
_atom_site.label_seq_id
_atom_site.Cartn_x
_atom_site.Cartn_y
_atom_site.Cartn_z
_atom_site.occupancy
_atom_site.B_iso_or_equiv
ATOM 1 C . CA GLY A 1 10.000 12.000 14.000 1.00 20.00
ATOM 2 C . CB GLY A 1 12.000 14.000 16.000 1.00 20.00
)";

    const char *kPdb =
        "CRYST1   40.000   50.000   60.000  90.00  90.00  90.00 P 1           1\n"
        "ATOM      1  CA  GLY A   1      10.000  12.000  14.000  1.00 20.00           C\n"
        "ATOM      2  CB  GLY A   1      12.000  14.000  16.000  1.00 20.00           C\n"
        "END\n";

    // A synthetic "protein": carbons filling one asymmetric unit of a small P2(1)2(1)2(1) cell. The
    // space group matters - in P1 the origin is free in all three directions, so |F| does not change
    // when the whole content is translated and there is no translation to recover. No specimen is
    // involved; the positions come out of a fixed seed.
    std::string ClusterPdb() {
        std::string pdb = "CRYST1   30.000   34.000   38.000  90.00  90.00  90.00 P 21 21 21    4\n";
        std::mt19937 rng(20260902);
        std::uniform_real_distribution<double> x(2, 14), y(2, 16), z(2, 18);
        char line[96];
        for (int i = 1; i <= 150; i++) {
            std::snprintf(line, sizeof line,
                          "ATOM  %5d  C   UNK A   1    %8.3f%8.3f%8.3f  1.00 20.00           C\n",
                          i, x(rng), y(rng), z(rng));
            pdb += line;
        }
        return pdb + "END\n";
    }

    std::string WriteTemp(const std::string &name, const char *content) {
        std::ofstream f(name);
        f << content;
        f.close();
        return name;
    }
}

// --model used to call gemmi::read_pdb unconditionally, so a deposited model handed over as mmCIF -
// which is how the PDB serves coordinates by default - was refused outright. Both formats now go
// through the same reader, which decides on the file's content.
TEST_CASE("ModelValidation_ReadsPdbAndMmcif", "[ModelValidation]") {
    Logger logger("ModelValidation_ReadsPdbAndMmcif");

    const auto pdb = WriteTemp("model_validation_test.pdb", kPdb);
    const auto cif = WriteTemp("model_validation_test.cif", kMmcif);

    const auto from_pdb = ModelReferenceIntensities(pdb, {}, {}, 4.0, logger);
    const auto from_cif = ModelReferenceIntensities(cif, {}, {}, 4.0, logger);

    REQUIRE_FALSE(from_pdb.empty());
    REQUIRE_FALSE(from_cif.empty());
    // The same structure either way, so the same reflections come out of it.
    CHECK(from_cif.size() == from_pdb.size());

    // The extension is not what decides: the same mmCIF under a .pdb name still reads.
    const auto misnamed = WriteTemp("model_validation_test_misnamed.pdb", kMmcif);
    CHECK_FALSE(ModelReferenceIntensities(misnamed, {}, {}, 4.0, logger).empty());

    std::filesystem::remove(pdb);
    std::filesystem::remove(cif);
    std::filesystem::remove(misnamed);
}

// A model that cannot be used must say why. Returning an empty result and logging was enough to lose
// the fact entirely: the run finished successfully with no R-free and no maps, which is exactly what
// a run without --model looks like.
TEST_CASE("ModelValidation_UnusableModelGivesAReason", "[ModelValidation]") {
    Logger logger("ModelValidation_UnusableModelGivesAReason");

    const auto empty_file = WriteTemp("model_validation_test_bogus.pdb", "not a coordinate file\n");

    const auto result = ValidateAgainstModel({}, UnitCell{.a = 40, .b = 50, .c = 60,
                                                          .alpha = 90, .beta = 90, .gamma = 90},
                                             empty_file, "", logger);
    CHECK_FALSE(result.ok);
    CHECK_FALSE(result.failure_reason.empty());
    CHECK_THAT(result.failure_reason, Catch::Matchers::ContainsSubstring(empty_file));

    const auto missing = ValidateAgainstModel({}, UnitCell{.a = 40, .b = 50, .c = 60,
                                                           .alpha = 90, .beta = 90, .gamma = 90},
                                              "model_validation_test_does_not_exist.pdb", "", logger);
    CHECK_FALSE(missing.ok);
    CHECK_FALSE(missing.failure_reason.empty());

    std::filesystem::remove(empty_file);
}

// A rigid-body step is only worth taking if it can find a shift it was not told about. The check is
// closed: the "observed" amplitudes are the model's own, so the answer is known to be zero shift.
TEST_CASE("ModelValidation_RigidBodyRecoversASmallShift", "[ModelValidation]") {
    Logger logger("ModelValidation_RigidBodyRecoversASmallShift");

    const auto path = WriteTemp("rigid_body_test.pdb", ClusterPdb().c_str());
    gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect);
    const gemmi::SpaceGroup *sg = st.find_spacegroup();
    REQUIRE(sg != nullptr);
    st.setup_cell_images();

    // "Observed" amplitudes: the model's own structure factors, so the target's minimum is exactly
    // where the model started.
    const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
    REQUIRE_FALSE(ref.empty());
    gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
    fobs.unit_cell_ = st.cell;
    fobs.spacegroup_ = sg;
    for (const auto &r : ref)
        fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
    fobs.ensure_sorted();

    const std::vector<gemmi::Position> original = ModelPositions(st.models[0]);
    std::vector<gemmi::Position> displaced;
    for (const gemmi::Position &p : original)
        displaced.emplace_back(p.x + 0.40, p.y - 0.30, p.z + 0.20);   // 0.54 A off
    SetModelPositions(st.models[0], displaced);

    const RigidBodyRefineResult result =
        RefineRigidBody(st.models[0], st.cell, *sg, fobs, 3.0, logger);
    CHECK(result.converged);

    const std::vector<gemmi::Position> refined = ModelPositions(st.models[0]);
    double before = 0, after = 0;
    for (size_t i = 0; i < original.size(); i++) {
        before += original[i].dist_sq(displaced[i]);
        after += original[i].dist_sq(refined[i]);
    }
    before = std::sqrt(before / original.size());
    after = std::sqrt(after / original.size());
    logger.Info("Rigid-body test: rmsd from the truth {:.3f} A -> {:.3f} A", before, after);
    CHECK(after < 0.2 * before);

    std::filesystem::remove(path);
}

// sigma_A is what says how much of the model to believe, so the two ends of its range are what the
// weighting has to get right: a model that explains the data completely, and one that explains none
// of it.
TEST_CASE("ModelValidation_SigmaAWeightsFollowTheModelsAgreement", "[ModelValidation]") {
    gemmi::UnitCell cell(40, 50, 60, 90, 90, 90);
    std::mt19937 rng(12345);
    std::normal_distribution<double> normal(0.0, 1.0);

    auto weights = [&](bool agreeing) {
        std::vector<SigmaAReflection> refl;
        for (int i = 0; i < 2000; i++) {
            SigmaAReflection r;
            r.f_calc = std::fabs(normal(rng)) * 100;
            r.f_obs = agreeing ? r.f_calc : std::fabs(normal(rng)) * 100;
            r.inv_d2 = 0.01 + 0.2 * (i / 2000.0);
            r.free = (i % 20) == 0;                 // the usual 5 %
            refl.push_back(r);
        }
        return EstimateSigmaA(refl, cell);
    };

    const SigmaAResult perfect = weights(true);
    const SigmaAResult useless = weights(false);
    CHECK(perfect.mean_fom > 0.85);
    CHECK(useless.mean_fom < 0.2);
    CHECK(perfect.shells == 2);   // 100 free reflections, 50 to a shell

    // No free reflections to estimate on: the coefficients are left alone rather than weighted by a
    // number that was never measured.
    std::vector<SigmaAReflection> no_free;
    for (int i = 0; i < 100; i++)
        no_free.push_back({100.0, 100.0, 0.05, 1, false, false});
    const SigmaAResult unweighted = EstimateSigmaA(no_free, cell);
    CHECK(unweighted.weight.size() == no_free.size());
    CHECK(unweighted.weight[0].m == 1.0);
    CHECK(unweighted.weight[0].d == 1.0);
}