// 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 files -> 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";
    WriteMtzReflections(reflections, TEST_CELL, experiment, provenance, original + ".mtz");

    MapCoefficientsMessage maps;
    maps.unit_cell = TEST_CELL;
    maps.space_group_name = "P 43 21 2";
    maps.wavelength_A = experiment.GetWavelength_A();
    for (size_t i = 0; i < reflections.size(); i++) {
        const float phi = static_cast<float>((i * 37) % 360);
        maps.h.push_back(reflections[i].h);
        maps.k.push_back(reflections[i].k);
        maps.l.push_back(reflections[i].l);
        maps.fp.push_back(reflections[i].F);
        maps.fc.push_back(reflections[i].F * 0.9f);
        maps.phic.push_back(phi);
        maps.fwt.push_back(reflections[i].F * 1.1f);
        maps.phwt.push_back(phi);
        maps.delfwt.push_back(reflections[i].F * 0.1f);
        maps.phdelwt.push_back(phi);
        maps.fom.push_back(0.8f);
        maps.free_flag.push_back(reflections[i].rfree_flag ? 0 : 1);
    }
    WriteMapCoefficientsMtz(maps, original + "_maps.mtz");

    WritePlacedModel(TestModel(), TEST_CELL, experiment.GetSpaceGroupOrP1(), original, logger);

    EndMessage end{};
    end.max_image_number = 100;
    FillRugnuxResultBlocks(original, end);
    REQUIRE(end.merged);
    REQUIRE(end.map_coefficients);
    REQUIRE(end.placed_model);
    CHECK(end.merged->h.size() == reflections.size());
    CHECK(end.merged->i_plus.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] == static_cast<float>(maps.k[5]));
    CHECK(maps_mtz.get_column_with_label("PHWT")[5] == maps.phwt[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{};
    end.merged = MergedReflectionsMessage{};
    end.merged->unit_cell = TEST_CELL;
    end.merged->space_group_name = "P 43 21 2";
    end.merged->h = {1}; end.merged->k = {2}; end.merged->l = {3};
    end.merged->i_mean = {10}; end.merged->sigma_i_mean = {1};
    end.merged->f = {3}; end.merged->sigma_f = {0.2};
    end.merged->free_flag = {1};

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