// 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 <random>
#include <string>
#include <vector>

#include <gemmi/dencalc.hpp>
#include <gemmi/fourier.hpp>
#include <gemmi/it92.hpp>
#include <gemmi/mtz.hpp>

#include "../common/Logger.h"
#include "../rugnux/ResultReport.h"
#include "../rugnux/RugnuxRefine.h"
#include "../writer/RugnuxReflectionFile.h"

namespace {

// A small synthetic orthorhombic cell (no real sample) and its reflections h, k, l = 1..6.
const gemmi::UnitCell test_cell(31.0, 37.0, 43.0, 90.0, 90.0, 90.0);

// An MTZ of `columns` (label, type), every row the hkl followed by values from `value(h, k, l, column)`.
gemmi::Mtz MakeMtz(const std::vector<std::pair<std::string, char>> &columns,
                   float (*value)(int h, int k, int l, size_t column)) {
    gemmi::Mtz mtz(true);
    mtz.spacegroup = gemmi::find_spacegroup_by_name("P 21 21 21");
    mtz.set_cell_for_all(test_cell);
    mtz.add_dataset("test").wavelength = 0.9;
    for (const auto &[label, type] : columns)
        mtz.add_column(label, type, -1, -1, false);
    for (int h = 1; h <= 6; h++)
        for (int k = 1; k <= 6; k++)
            for (int l = 1; l <= 6; l++) {
                mtz.data.insert(mtz.data.end(), {float(h), float(k), float(l)});
                for (size_t c = 0; c < columns.size(); c++)
                    mtz.data.push_back(value(h, k, l, c));
                mtz.nreflections++;
            }
    return mtz;
}

bool InTestSet(int h, int k, int l) { return (h * 7 + k * 13 + l * 17) % 10 == 0; }

// Waters scattered over the test cell, and their own |F|^2 as the data: the model as given scores well.
// With `strict_ncs` the second half of the waters is the first half moved by a translation, and the model
// keeps only the first half with a strict-NCS operator (not given) that generates the rest.
void SyntheticRefineInput(gemmi::Structure &st, RugnuxReflections &data, bool strict_ncs = false) {
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    st.cell = test_cell;
    st.spacegroup_hm = sg->xhm();
    st.setup_cell_images();
    gemmi::Model model(1);
    gemmi::Chain chain("A");
    const gemmi::Position ncs_shift(5.3, 7.1, 2.9);
    std::mt19937 rng(7);
    std::uniform_real_distribution<double> frac(0.0, 1.0);
    for (int i = 0; i < 120; i++) {
        gemmi::Residue res;
        res.name = "HOH";
        res.seqid = gemmi::SeqId(i + 1, ' ');
        res.het_flag = 'H';
        gemmi::Atom a;
        a.name = "O";
        a.element = gemmi::El::O;
        a.pos = test_cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
        if (strict_ncs && i >= 60)
            a.pos = chain.residues[i - 60].atoms[0].pos + ncs_shift;
        a.b_iso = 15.0f;
        a.occ = 1.0f;
        res.atoms.push_back(a);
        chain.residues.push_back(res);
    }
    model.chains.push_back(chain);
    st.models.push_back(model);

    gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
    dc.d_min = 2.0;
    dc.grid.unit_cell = test_cell;
    dc.grid.spacegroup = sg;
    dc.set_refmac_compatible_blur(st.models[0]);
    dc.put_model_density_on_grid(st.models[0]);
    const gemmi::AsuData<std::complex<float>> f =
            gemmi::transform_map_to_f_phi(dc.grid, true).prepare_asu_data(2.0, dc.blur);

    data.name = "synthetic.mtz";
    data.source = "MTZ";
    data.cell = test_cell;
    data.space_group = sg->xhm();
    data.amplitudes = false;
    for (const auto &hv : f.v) {
        if (hv.hkl == gemmi::Miller{{0, 0, 0}})
            continue;
        const double fc = std::abs(hv.value);
        MergedReflection m;
        m.h = hv.hkl[0];
        m.k = hv.hkl[1];
        m.l = hv.hkl[2];
        m.d = static_cast<float>(test_cell.calculate_d(hv.hkl));
        m.I = static_cast<float>(fc * fc);
        m.sigma = static_cast<float>(0.02 * fc * fc + 1.0);
        data.reflections.push_back(m);
    }
    if (strict_ncs) {
        st.models[0].chains[0].residues.resize(60);
        gemmi::Transform shift;
        shift.vec = ncs_shift;
        st.ncs.push_back({"1", true, gemmi::Transform()});
        st.ncs.push_back({"2", false, shift});
    }
}

}  // namespace

TEST_CASE("RugnuxReflectionsFromMtz takes observed amplitudes over map coefficients", "[refine_only]") {
    // FWT comes first and is an amplitude with a phase after it; FP/SIGFP is what is observed.
    const gemmi::Mtz mtz = MakeMtz({{"FWT", 'F'}, {"PHWT", 'P'}, {"FP", 'F'}, {"SIGFP", 'Q'}, {"FreeR_flag", 'I'}},
                                   [](int h, int k, int l, size_t c) -> float {
                                       switch (c) {
                                           case 0: return 999.0f;
                                           case 1: return 10.0f;
                                           case 2: return 100.0f + h;
                                           case 3: return 2.0f;
                                           default: return InTestSet(h, k, l) ? 0.0f : float(1 + (h + k + l) % 19);
                                       }
                                   });
    const RugnuxReflections r = RugnuxReflectionsFromMtz(mtz, "synthetic.mtz");
    CHECK(r.amplitudes);
    CHECK(r.value_column == "FP");
    CHECK(r.sigma_column == "SIGFP");
    CHECK(r.free_column == "FreeR_flag");
    CHECK(r.wavelength_A == Catch::Approx(0.9));
    REQUIRE(r.reflections.size() == 216);
    int n_free = 0;
    for (const MergedReflection &m : r.reflections) {
        CHECK(m.F == Catch::Approx(100.0 + m.h));
        CHECK(m.rfree_flag == InTestSet(m.h, m.k, m.l));
        n_free += m.rfree_flag;
    }
    CHECK(r.n_free == n_free);
}

TEST_CASE("RugnuxReflectionsFromMtz falls back to intensities and reads a 1-marks-free flag", "[refine_only]") {
    const gemmi::Mtz mtz = MakeMtz({{"IMEAN", 'J'}, {"SIGIMEAN", 'Q'}, {"R-free-flags", 'I'}},
                                   [](int h, int k, int l, size_t c) -> float {
                                       if (c == 0) return 50.0f * h;
                                       if (c == 1) return 5.0f;
                                       return InTestSet(h, k, l) ? 1.0f : 0.0f;
                                   });
    const RugnuxReflections r = RugnuxReflectionsFromMtz(mtz, "synthetic.mtz");
    CHECK_FALSE(r.amplitudes);
    CHECK(r.value_column == "IMEAN");
    CHECK(r.free_column == "R-free-flags");
    for (const MergedReflection &m : r.reflections) {
        CHECK(m.I == Catch::Approx(50.0 * m.h));
        CHECK(std::isnan(m.F));
        CHECK(m.rfree_flag == InTestSet(m.h, m.k, m.l));
    }
}

TEST_CASE("RugnuxReflectionsFromMtz without free flags, and without any value column", "[refine_only]") {
    const gemmi::Mtz plain = MakeMtz({{"F", 'F'}, {"SIGF", 'Q'}},
                                     [](int h, int, int, size_t c) -> float { return c == 0 ? 10.0f * h : 1.0f; });
    const RugnuxReflections r = RugnuxReflectionsFromMtz(plain, "synthetic.mtz");
    CHECK(r.value_column == "F");
    CHECK(r.free_column.empty());
    CHECK(r.n_free == 0);

    const gemmi::Mtz calculated = MakeMtz({{"FC", 'F'}, {"PHIC", 'P'}},
                                          [](int, int, int, size_t) -> float { return 1.0f; });
    CHECK_THROWS(RugnuxReflectionsFromMtz(calculated, "synthetic.mtz"));
}

TEST_CASE("ParseRugnuxReflections reads an MTZ from its bytes", "[refine_only]") {
    const gemmi::Mtz mtz = MakeMtz({{"F", 'F'}, {"SIGF", 'Q'}, {"FreeR_flag", 'I'}},
                                   [](int h, int k, int l, size_t c) -> float {
                                       if (c == 0) return 10.0f * h;
                                       if (c == 1) return 1.0f;
                                       return InTestSet(h, k, l) ? 0.0f : 1.0f;
                                   });
    std::string bytes;
    mtz.write_to_string(bytes);
    const RugnuxReflections r = ParseRugnuxReflections(bytes, "synthetic.mtz");
    CHECK(r.source == "MTZ");
    CHECK(r.space_group == "P 21 21 21");
    CHECK(r.reflections.size() == 216);
    CHECK(r.free_column == "FreeR_flag");
}

TEST_CASE("RunRefineOnly scores a model against intensities, with amplitudes and free flags made for it",
          "[refine_only]") {
    gemmi::Structure st;
    RugnuxReflections data;
    SyntheticRefineInput(st, data);
    RefinementSettings settings;
    settings.score_only = true;
    settings.gpu = false;
    settings.nthreads = 2;
    Logger logger("RefineOnlyTest");
    const RefineOnlyResult r = RunRefineOnly(data, st, 2.5, settings, logger);
    REQUIRE(r.refinement.ok);
    CHECK(r.french_wilson);
    CHECK(r.free_flags_generated);
    CHECK(r.free_reflections > 0);
    CHECK(r.d_min_limit == 2.5);
    CHECK(r.refinement.d_min >= 2.5);
    CHECK(r.refinement.r_work_start < 0.15);
    CHECK(r.refinement.score_start.r_work == Catch::Approx(r.refinement.r_work_start));
    CHECK(r.refinement.score_start.shells.size() == 20);
    CHECK_FALSE(r.refinement.restrained);

    const ReportDocument doc = BuildRefineReportDocument("out", "synthetic.mtz", "model.cif", r);
    CHECK(doc.verdict == "OK");
    const std::string text = RenderReportText(doc, true);
    CHECK(text.find("REFINE= PERFORMED") != std::string::npos);
    CHECK(text.find("REFINE_SCORE_R_WORK=") != std::string::npos);
    CHECK(text.find("REFINE_FREE_FLAGS= GENERATED") != std::string::npos);
    CHECK(text.find("REFINE_FRENCH_WILSON= TRUE") != std::string::npos);
}

TEST_CASE("ParseRugnuxReflections keeps a structure-factor mmCIF's free-flag classes beside its status", "[refine_only]") {
    // status marks class 5 free; the classes are h % 3 (three of them)
    std::string cif = "data_r\n_cell.length_a 31\n_cell.length_b 37\n_cell.length_c 43\n_cell.angle_alpha 90\n"
                      "_cell.angle_beta 90\n_cell.angle_gamma 90\n_symmetry.space_group_name_H-M 'P 21 21 21'\n"
                      "loop_\n_refln.index_h\n_refln.index_k\n_refln.index_l\n_refln.status\n"
                      "_refln.pdbx_r_free_flag\n_refln.F_meas_au\n_refln.F_meas_sigma_au\n";
    for (int h = 1; h <= 6; h++)
        for (int k = 1; k <= 3; k++)
            cif += fmt::format("{} {} 1 {} {} {} 1\n", h, k, h % 3 == 2 ? 'f' : 'o', h % 3 == 2 ? 5 : h % 3, 10 * h);
    const RugnuxReflections r = ParseRugnuxReflections(cif, "synthetic-sf.cif");
    REQUIRE(r.reflections.size() == 18);
    REQUIRE(r.free_class.size() == 18);
    CHECK(r.n_free == 6);
    for (size_t i = 0; i < r.reflections.size(); i++) {
        CHECK(r.reflections[i].rfree_flag == (r.reflections[i].h % 3 == 2));
        CHECK(r.free_class[i] == (r.reflections[i].h % 3 == 2 ? 5 : r.reflections[i].h % 3));
    }
}

namespace {

// The synthetic input with its reflections in 20 free-flag classes, class 5 marked free, and the intensities of
// class `held_out` disturbed by 30 % at random: the class a model was not refined against, which it fits worst.
// `count`: the reflections of that class.
void ClassedRefineInput(gemmi::Structure &st, RugnuxReflections &data, int held_out, int &count) {
    SyntheticRefineInput(st, data);
    std::mt19937 rng(11);
    std::uniform_real_distribution<float> noise(0.7f, 1.3f);
    data.free_column = "FreeR_flag";
    count = 0;
    for (MergedReflection &m : data.reflections) {
        const int c = (m.h * 7 + m.k * 13 + m.l * 17 + 1000) % 20;
        data.free_class.push_back(c);
        m.rfree_flag = c == 5;
        data.n_free += m.rfree_flag;
        if (c == held_out) {
            m.I *= noise(rng) * noise(rng);
            count++;
        }
    }
}

}  // namespace

TEST_CASE("RunRefineOnly takes the free-flag class the model fits worst, by the model file's test-set count",
          "[refine_only]") {
    gemmi::Structure st;
    RugnuxReflections data;
    int count = 0;
    ClassedRefineInput(st, data, 3, count);
    st.meta.refinement.emplace_back();
    st.meta.refinement[0].rfree_set_count = count;
    RefinementSettings settings;
    settings.score_only = true;
    settings.gpu = false;
    settings.nthreads = 2;
    Logger logger("RefineOnlyTest");
    const RefineOnlyResult r = RunRefineOnly(data, st, 0, settings, logger);
    REQUIRE(r.refinement.ok);
    CHECK(r.free_class == 3);
    CHECK(r.free_set_source == "REMARK3_COUNT");
    CHECK(r.free_reflections == count);
    const std::string text = RenderReportText(BuildRefineReportDocument("out", "x-sf.cif", "model.cif", r), true);
    CHECK(text.find("REFINE_FREE_SET_SOURCE= REMARK3_COUNT") != std::string::npos);
    CHECK(text.find("REFINE_FREE_SET_CLASS= 3") != std::string::npos);
}

TEST_CASE("RunRefineOnly takes the worst-fit free-flag class where the model file gives no count", "[refine_only]") {
    gemmi::Structure st;
    RugnuxReflections data;
    int count = 0;
    ClassedRefineInput(st, data, 3, count);
    RefinementSettings settings;
    settings.score_only = true;
    settings.gpu = false;
    settings.nthreads = 2;
    Logger logger("RefineOnlyTest");
    const RefineOnlyResult r = RunRefineOnly(data, st, 0, settings, logger);
    REQUIRE(r.refinement.ok);
    CHECK(r.free_class == 3);
    CHECK(r.free_set_source == "WORST_FIT");
}

TEST_CASE("RunRefineOnly keeps the marked free-flag class where the model fits it worst", "[refine_only]") {
    gemmi::Structure st;
    RugnuxReflections data;
    int count = 0;
    ClassedRefineInput(st, data, 5, count);
    RefinementSettings settings;
    settings.score_only = true;
    settings.gpu = false;
    settings.nthreads = 2;
    Logger logger("RefineOnlyTest");
    const RefineOnlyResult r = RunRefineOnly(data, st, 0, settings, logger);
    REQUIRE(r.refinement.ok);
    CHECK(r.free_class == 5);
    CHECK(r.free_set_source == "STATUS");
    CHECK(r.free_reflections == count);
}

TEST_CASE("RunRefineOnly scores a model with strict-NCS operators as the expanded model", "[refine_only]") {
    gemmi::Structure st;
    RugnuxReflections data;
    SyntheticRefineInput(st, data, true);
    RefinementSettings settings;
    settings.score_only = true;
    settings.gpu = false;
    settings.nthreads = 2;
    Logger logger("RefineOnlyTest");
    const RefineOnlyResult r = RunRefineOnly(data, st, 2.5, settings, logger);
    REQUIRE(r.refinement.ok);
    CHECK(r.refinement.r_work_start < 0.15);
}

TEST_CASE("RunRefineOnly reports a refinement that throws as not performed, with the reason", "[refine_only]") {
    // The data's cell, a hundred metres on a side, makes the bulk-solvent mask's grid too large to allocate.
    gemmi::Structure st;
    RugnuxReflections data;
    SyntheticRefineInput(st, data);
    data.cell = gemmi::UnitCell(1e6, 1e6, 1e6, 90, 90, 90);
    RefinementSettings settings;
    settings.gpu = false;
    settings.nthreads = 2;
    Logger logger("RefineOnlyTest");
    RefineOnlyResult r;
    REQUIRE_NOTHROW(r = RunRefineOnly(data, st, 0.0, settings, logger));
    CHECK_FALSE(r.refinement.ok);
    CHECK(r.refinement.failure_reason.find("refinement failed") == 0);
    const std::string text = RenderReportText(BuildRefineReportDocument("out", "synthetic.mtz", "model.cif", r), true);
    CHECK(text.find("REFINE= NOT_PERFORMED") != std::string::npos);
    CHECK(text.find("REFINE_REASON= refinement failed") != std::string::npos);
}