// 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 <fstream>
#include <iostream>
#include <iterator>

#include <gemmi/mtz.hpp>
#include <gemmi/mmread_gz.hpp>
#include <gemmi/calculate.hpp>  // count_atom_sites

#include "../common/Definitions.h"
#include "../common/DiffractionExperiment.h"
#include "../common/Logger.h"
#include "../image_analysis/WriteReflections.h"
#include "../frame_serialize/CBORStream2Serializer.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../writer/WriteModel.h"
#include "../writer/WriteRugnuxFiles.h"
#include "../writer/RugnuxResultBlocks.h"
#include "SyntheticMergedReflections.h"

namespace {
    constexpr UnitCell TEST_CELL{47.0f, 47.0f, 63.0f, 90.0f, 90.0f, 90.0f};

    std::string ReadFile(const std::string &path) {
        std::ifstream f(path, std::ios::binary);
        return {std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()};
    }

    std::vector<std::string> Labels(const gemmi::Mtz &mtz) {
        std::vector<std::string> v;
        for (const auto &c : mtz.columns)
            v.push_back(c.label);
        return v;
    }

    // A small model with what a placed model carries: two chains, a ligand, a water, an alternate
    // conformation, an anisotropic atom, an insertion code and a charge.
    gemmi::Structure TestModel() {
        gemmi::Structure st;
        st.name = "test";
        st.cell = TEST_CELL;
        st.spacegroup_hm = "P 43 21 2";
        gemmi::Model &m = st.models.emplace_back(1);
        double x = 1.0;
        auto add_atom = [&x](gemmi::Residue &r, const std::string &name, const char *el, char altloc = 0) {
            gemmi::Atom &a = r.atoms.emplace_back();
            a.name = name;
            a.element = gemmi::Element(el);
            a.altloc = altloc;
            a.pos = gemmi::Position(x, 2.0 * x + 0.123, 3.0 * x - 0.456);
            a.b_iso = 15.0f + static_cast<float>(x);
            a.occ = altloc ? 0.5f : 1.0f;
            x += 1.37;
            return &a;
        };
        for (const std::string chain : {"A", "B"}) {
            gemmi::Chain &ch = m.chains.emplace_back(chain);
            for (int i = 1; i <= 3; i++) {
                gemmi::Residue &r = ch.residues.emplace_back(gemmi::ResidueId{gemmi::SeqId(i, i == 3 ? 'A' : ' '), {}, "ALA"});
                r.het_flag = 'A';
                add_atom(r, "N", "N");
                add_atom(r, "CA", "C");
                add_atom(r, "C", "C");
                add_atom(r, "O", "O");
                add_atom(r, "CB", "C", 'A');
                add_atom(r, "CB", "C", 'B');
            }
        }
        gemmi::Residue &zn = m.chains[0].residues.emplace_back(gemmi::ResidueId{gemmi::SeqId(101, ' '), {}, "ZN"});
        zn.het_flag = 'H';
        gemmi::Atom *a = add_atom(zn, "ZN", "Zn");
        a->charge = 2;
        a->aniso = {0.2f, 0.25f, 0.3f, 0.01f, -0.02f, 0.03f};
        gemmi::Residue &w = m.chains[1].residues.emplace_back(gemmi::ResidueId{gemmi::SeqId(201, ' '), {}, "HOH"});
        w.het_flag = 'H';
        add_atom(w, "O", "O");
        st.setup_cell_images();
        return st;
    }
}

TEST_CASE("RugnuxResultFiles: rugnux's outputs -> END blocks -> CBOR -> writer, byte for byte", "[rugnux_result_files]") {
    Logger logger("RugnuxResultFiles");
    const std::filesystem::path dir = std::filesystem::temp_directory_path() / "rugnux_result_files_test";
    std::filesystem::remove_all(dir);
    std::filesystem::create_directories(dir);
    const std::string original = (dir / "rugnux").string();
    const std::string written = (dir / "writer").string();

    // About the size of a standard-crystal run (~25k reflections), with an anomalous split.
    jfjoch_test::SyntheticMergeParams params;
    params.true_space_group = "P 43 21 2";
    params.twin_supergroup = "P 43 21 2";
    params.d_min_A = 2.25;
    auto reflections = jfjoch_test::GenerateSyntheticMerged(params);
    for (size_t i = 0; i < reflections.size(); i++) {
        auto &r = reflections[i];
        r.I_plus = r.I * 1.01f; r.sigma_plus = r.sigma * 1.4f;
        r.F_plus = std::sqrt(std::max(r.I, 1.0f)); r.sigmaF_plus = 0.1f;
        if (i % 7 != 0) {
            r.I_minus = r.I * 0.99f; r.sigma_minus = r.sigma * 1.4f;
            r.F_minus = std::sqrt(std::max(r.I, 1.0f)); r.sigmaF_minus = 0.12f;
        }
        r.rfree_flag = (i % 20 == 0);
    }
    DiffractionExperiment experiment;
    experiment.IncidentEnergy_keV(12.7f);
    experiment.SpaceGroupNumber(96);
    experiment.SetUnitCell(TEST_CELL);
    OutputProvenance provenance;
    provenance.command_line = "rugnux -N 8 input_master.h5";
    provenance.input_file = "/some/where/input_master.h5";
    // The files as rugnux writes them from its outputs (RugnuxOutputFiles.cpp), and the same outputs as the
    // broker puts them in the END message (JFJochReceiverRugnux::AddToEndMessage).
    EndMessage end{};
    end.max_image_number = 100;
    end.merged = MergedMtzTable(reflections, TEST_CELL, experiment, provenance);
    WriteMtz(*end.merged, original + ".mtz");

    // Map coefficients as model validation lays them out (ModelMaps::coefficients).
    MtzTable maps;
    maps.unit_cell = TEST_CELL;
    maps.space_group_name = "P 43 21 2";
    maps.datasets = {{"HKL_base", "HKL_base", "HKL_base", 0.0f},
                     {"model_validation", "model_validation", "model_validation", experiment.GetWavelength_A()}};
    for (const auto &[label, type] : std::vector<std::pair<const char *, char>>{
             {"H", 'H'}, {"K", 'H'}, {"L", 'H'}, {"FP", 'F'}, {"FC", 'F'}, {"PHIC", 'P'}, {"FWT", 'F'},
             {"PHWT", 'P'}, {"DELFWT", 'F'}, {"PHDELWT", 'P'}, {"FOM", 'W'}, {"FREE", 'I'}})
        maps.columns.push_back({label, type, type == 'H' ? 0 : 1, {}});
    for (size_t i = 0; i < reflections.size(); i++) {
        const float phi = static_cast<float>((i * 37) % 360);
        const float row[] = {static_cast<float>(reflections[i].h), static_cast<float>(reflections[i].k),
                             static_cast<float>(reflections[i].l), reflections[i].F, reflections[i].F * 0.9f, phi,
                             reflections[i].F * 1.1f, phi, reflections[i].F * 0.1f, phi, 0.8f,
                             reflections[i].rfree_flag ? 0.0f : 1.0f};
        for (size_t c = 0; c < maps.columns.size(); c++)
            maps.columns[c].values.push_back(row[c]);
    }
    WriteMtz(maps, original + "_maps.mtz");
    end.map_coefficients = maps;

    WritePlacedModel(TestModel(), TEST_CELL, experiment.GetSpaceGroupOrP1(), original, logger);
    end.placed_model = PlacedModelFromStructure(TestModel(), TEST_CELL, experiment.GetSpaceGroupOrP1().xhm());
    REQUIRE(end.merged->columns.size() == 16);   // with the anomalous columns
    CHECK(end.merged->columns[0].values.size() == reflections.size());
    CHECK(end.placed_model->atom_name.size() == 2 * 3 * 6 + 2);

    std::vector<uint8_t> buffer(MESSAGE_SIZE_FOR_START_END);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());
    REQUIRE_NOTHROW(serializer.SerializeSequenceEnd(end));
    std::cout << "END message with " << reflections.size() << " merged reflections, as many map coefficients and "
              << end.placed_model->atom_name.size() << " atoms: " << serializer.GetBufferSize() << " bytes" << std::endl;
    auto deserialized = CBORStream2Deserialize(buffer.data(), serializer.GetBufferSize());
    REQUIRE(deserialized);
    REQUIRE(deserialized->end_message);
    const EndMessage &out = *deserialized->end_message;
    REQUIRE(out.merged);
    REQUIRE(out.map_coefficients);
    REQUIRE(out.placed_model);
    CHECK(out.merged->history == end.merged->history);
    CHECK(out.placed_model->x == end.placed_model->x);

    WriteRugnuxResultFiles(out, written, false, logger);

    CHECK(ReadFile(written + ".mtz") == ReadFile(original + ".mtz"));
    CHECK(ReadFile(written + "_maps.mtz") == ReadFile(original + "_maps.mtz"));
    CHECK(ReadFile(written + "_model.cif") == ReadFile(original + "_model.cif"));
    CHECK(ReadFile(written + "_model.pdb") == ReadFile(original + "_model.pdb"));

    const gemmi::Mtz mtz = gemmi::read_mtz_file(written + ".mtz");
    CHECK(Labels(mtz) == std::vector<std::string>{"H", "K", "L", "IMEAN", "SIGIMEAN", "I(+)", "SIGI(+)", "I(-)",
                                                  "SIGI(-)", "F", "SIGF", "F(+)", "SIGF(+)", "F(-)", "SIGF(-)",
                                                  "FreeR_flag"});
    CHECK(mtz.nreflections == static_cast<int>(reflections.size()));
    CHECK(mtz.datasets.back().wavelength == Catch::Approx(experiment.GetWavelength_A()).epsilon(1e-5));
    const gemmi::Mtz maps_mtz = gemmi::read_mtz_file(written + "_maps.mtz");
    CHECK(Labels(maps_mtz) == std::vector<std::string>{"H", "K", "L", "FP", "FC", "PHIC", "FWT", "PHWT", "DELFWT",
                                                       "PHDELWT", "FOM", "FREE"});
    CHECK(maps_mtz.get_column_with_label("K")[5] == maps.columns[1].values[5]);
    CHECK(maps_mtz.get_column_with_label("PHWT")[5] == maps.columns[7].values[5]);
    const gemmi::Structure st = gemmi::read_structure_gz(written + "_model.cif");
    CHECK(gemmi::count_atom_sites(st) == 2 * 3 * 6 + 2);

    std::filesystem::remove_all(dir);
}

TEST_CASE("RugnuxResultFiles: an existing file is kept when overwrite is off", "[rugnux_result_files]") {
    Logger logger("RugnuxResultFiles");
    const std::filesystem::path dir = std::filesystem::temp_directory_path() / "rugnux_result_files_overwrite";
    std::filesystem::remove_all(dir);
    std::filesystem::create_directories(dir);
    const std::string prefix = (dir / "run").string();

    EndMessage end{};
    DiffractionExperiment experiment;
    experiment.SpaceGroupNumber(96);
    end.merged = MergedMtzTable({MergedReflection{.h = 1, .k = 2, .l = 3, .I = 10, .sigma = 1, .F = 3, .sigmaF = 0.2}},
                                TEST_CELL, experiment, OutputProvenance{});

    { std::ofstream f(prefix + ".mtz"); f << "keep"; }
    WriteRugnuxResultFiles(end, prefix, false, logger);
    CHECK(ReadFile(prefix + ".mtz") == "keep");
    CHECK(!std::filesystem::exists(prefix + ".mtz.tmp"));

    WriteRugnuxResultFiles(end, prefix, true, logger);
    const gemmi::Mtz mtz = gemmi::read_mtz_file(prefix + ".mtz");
    CHECK(mtz.nreflections == 1);
    CHECK(Labels(mtz) == std::vector<std::string>{"H", "K", "L", "IMEAN", "SIGIMEAN", "F", "SIGF", "FreeR_flag"});

    std::filesystem::remove_all(dir);
}

TEST_CASE("RugnuxResultFiles: an MTZ table that describes no file is refused", "[rugnux_result_files]") {
    const std::string path = (std::filesystem::temp_directory_path() / "rugnux_result_files_refused.mtz").string();
    DiffractionExperiment experiment;
    experiment.SpaceGroupNumber(96);
    const MtzTable good = MergedMtzTable({MergedReflection{.h = 1, .k = 2, .l = 3, .I = 10, .sigma = 1, .F = 3,
                                                           .sigmaF = 0.2}},
                                         TEST_CELL, experiment, OutputProvenance{});
    REQUIRE_NOTHROW(WriteMtz(good, path));

    MtzTable short_column = good;
    short_column.columns.back().values.clear();
    CHECK_THROWS(WriteMtz(short_column, path));

    MtzTable no_dataset = good;
    no_dataset.columns.back().dataset_id = 2;
    CHECK_THROWS(WriteMtz(no_dataset, path));

    MtzTable unknown_group = good;
    unknown_group.space_group_name = "not a group";
    CHECK_THROWS(WriteMtz(unknown_group, path));

    MtzTable bad_batch = good;
    bad_batch.batches.push_back({1, "Rugnux", std::vector<int32_t>(28, 0), std::vector<float>(156, 0.0f), {}});
    CHECK_THROWS(WriteMtz(bad_batch, path));

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

TEST_CASE("RugnuxResultFiles: a placed model block that describes no model is refused", "[rugnux_result_files]") {
    Logger logger("RugnuxResultFiles");
    const PlacedModelMessage good = PlacedModelFromStructure(TestModel(), TEST_CELL, "P 43 21 2");
    REQUIRE_NOTHROW(PlacedModelStructure(good));

    PlacedModelMessage short_field = good;
    short_field.b_iso.pop_back();
    CHECK_THROWS(PlacedModelStructure(short_field));

    PlacedModelMessage bad_aniso = good;
    bad_aniso.aniso.resize(5);
    CHECK_THROWS(PlacedModelStructure(bad_aniso));

    // The writer reports it and goes on: no file, no crash.
    const std::filesystem::path dir = std::filesystem::temp_directory_path() / "rugnux_result_files_bad_model";
    std::filesystem::remove_all(dir);
    std::filesystem::create_directories(dir);
    const std::string prefix = (dir / "run").string();
    EndMessage end{};
    end.placed_model = short_field;
    REQUIRE_NOTHROW(WriteRugnuxResultFiles(end, prefix, false, logger));
    end.placed_model = good;
    end.placed_model->space_group_name = "not a group";
    REQUIRE_NOTHROW(WriteRugnuxResultFiles(end, prefix, false, logger));
    CHECK(!std::filesystem::exists(prefix + "_model.cif"));
    std::filesystem::remove_all(dir);
}