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

}  // 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]") {
    // Atoms scattered over the cell; the data are their own |F|^2, so the model as given scores well.
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    gemmi::Structure st;
    st.cell = test_cell;
    st.spacegroup_hm = sg->xhm();
    st.setup_cell_images();
    gemmi::Model model(1);
    gemmi::Chain chain("A");
    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)));
        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);

    RugnuxReflections data;
    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);
    }

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