// 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 <filesystem>
#include <map>
#include <tuple>
#include <vector>

#include <gemmi/mtz.hpp>

#include "../image_analysis/LoadFCalcFromMtz.h"

namespace {
    // Write a small P4(3)2(1)2 MTZ with an IMEAN column and a FreeR_flag column whose free reflections
    // carry `free_value` (and the rest its complement). ~1 in 10 reflections is flagged free.
    std::string WriteMtzWithFreeR(const std::string& path, int free_value) {
        gemmi::Mtz mtz(true);
        mtz.spacegroup = gemmi::find_spacegroup_by_number(96);
        mtz.set_cell_for_all(gemmi::UnitCell(78.0, 78.0, 37.0, 90.0, 90.0, 90.0));
        mtz.add_dataset("test");
        mtz.add_column("IMEAN", 'J', -1, -1, false);
        mtz.add_column("FreeR_flag", 'I', -1, -1, false);
        std::vector<float> data;
        int nref = 0;
        for (int h = 1; h <= 10; ++h)
            for (int k = 1; k <= 10; ++k)
                for (int l = 1; l <= 6; ++l) {
                    const bool is_free = ((h * 7 + k * 13 + l * 17) % 10 == 0);
                    const int flag = is_free ? free_value : (free_value == 0 ? 1 : 0);
                    data.insert(data.end(), {static_cast<float>(h), static_cast<float>(k),
                                             static_cast<float>(l), 100.0f, static_cast<float>(flag)});
                    ++nref;
                }
        mtz.nreflections = nref;
        mtz.data = std::move(data);
        mtz.write_to_file(path);
        return path;
    }

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

TEST_CASE("LoadReferenceMtz imports FreeR (CCP4 convention: test set = 0)", "[reference_mtz]") {
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_freer_ccp4.mtz").string();
    WriteMtzWithFreeR(path, /*free_value=*/0);
    const auto ref = LoadReferenceMtz(path);
    std::filesystem::remove(path);

    REQUIRE(ref.has_free_flags);
    CHECK(ref.n_free > 0);
    CHECK(ref.n_free < static_cast<int>(ref.reflections.size()) / 2);   // free is the minority
    for (const auto& r : ref.reflections)
        CHECK(r.rfree_flag == ExpectedFree(r.h, r.k, r.l));
}

TEST_CASE("LoadReferenceMtz imports FreeR (phenix convention: test set = 1) via the complement", "[reference_mtz]") {
    // Here flag 0 is the majority (work), so the loader must take the complement and treat 1 as free.
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_freer_phenix.mtz").string();
    WriteMtzWithFreeR(path, /*free_value=*/1);
    const auto ref = LoadReferenceMtz(path);
    std::filesystem::remove(path);

    REQUIRE(ref.has_free_flags);
    CHECK(ref.n_free > 0);
    CHECK(ref.n_free < static_cast<int>(ref.reflections.size()) / 2);
    for (const auto& r : ref.reflections)
        CHECK(r.rfree_flag == ExpectedFree(r.h, r.k, r.l));
}

TEST_CASE("LoadReferenceMtz without a FreeR column reports no free flags", "[reference_mtz]") {
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_no_freer.mtz").string();
    gemmi::Mtz mtz(true);
    mtz.spacegroup = gemmi::find_spacegroup_by_number(96);
    mtz.set_cell_for_all(gemmi::UnitCell(78.0, 78.0, 37.0, 90.0, 90.0, 90.0));
    mtz.add_dataset("test");
    mtz.add_column("IMEAN", 'J', -1, -1, false);
    std::vector<float> data;
    int nref = 0;
    for (int h = 1; h <= 6; ++h)
        for (int k = 1; k <= 6; ++k)
            for (int l = 1; l <= 4; ++l) {
                data.insert(data.end(), {static_cast<float>(h), static_cast<float>(k),
                                         static_cast<float>(l), 100.0f});
                ++nref;
            }
    mtz.nreflections = nref;
    mtz.data = std::move(data);
    mtz.write_to_file(path);
    const auto ref = LoadReferenceMtz(path);
    std::filesystem::remove(path);

    CHECK_FALSE(ref.has_free_flags);
    CHECK(ref.n_free == 0);
}