Jungfraujoch/writer/RugnuxResultBlocks.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "RugnuxResultBlocks.h"

#include <filesystem>

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

namespace {

UnitCell ToUnitCell(const gemmi::UnitCell &c) {
    return UnitCell{static_cast<float>(c.a), static_cast<float>(c.b), static_cast<float>(c.c),
                    static_cast<float>(c.alpha), static_cast<float>(c.beta), static_cast<float>(c.gamma)};
}

// The column with this label, empty where the file has none.
std::vector<float> Column(const gemmi::Mtz &mtz, const std::string &label) {
    std::vector<float> v;
    if (const gemmi::Mtz::Column *col = mtz.column_with_label(label))
        for (int i = 0; i < col->size(); i++)
            v.push_back((*col)[i]);
    return v;
}

std::vector<int32_t> IntColumn(const gemmi::Mtz &mtz, const std::string &label) {
    std::vector<int32_t> v;
    for (float x : Column(mtz, label))
        v.push_back(static_cast<int32_t>(x));
    return v;
}

std::vector<uint8_t> FlagColumn(const gemmi::Mtz &mtz, const std::string &label) {
    std::vector<uint8_t> v;
    for (float x : Column(mtz, label))
        v.push_back(static_cast<uint8_t>(x));
    return v;
}

std::string SpaceGroupName(const gemmi::Mtz &mtz) {
    return mtz.spacegroup != nullptr ? mtz.spacegroup->xhm() : std::string();
}

}  // namespace

MergedReflectionsMessage MergedReflectionsFromMtz(const std::string &path) {
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    MergedReflectionsMessage m;
    m.unit_cell = ToUnitCell(mtz.cell);
    m.space_group_name = SpaceGroupName(mtz);
    m.wavelength_A = static_cast<float>(mtz.datasets.back().wavelength);
    m.crystal_name = mtz.datasets.back().crystal_name;
    m.history = mtz.history;
    m.h = IntColumn(mtz, "H");
    m.k = IntColumn(mtz, "K");
    m.l = IntColumn(mtz, "L");
    m.i_mean = Column(mtz, "IMEAN");
    m.sigma_i_mean = Column(mtz, "SIGIMEAN");
    m.i_plus = Column(mtz, "I(+)");
    m.sigma_i_plus = Column(mtz, "SIGI(+)");
    m.i_minus = Column(mtz, "I(-)");
    m.sigma_i_minus = Column(mtz, "SIGI(-)");
    m.f = Column(mtz, "F");
    m.sigma_f = Column(mtz, "SIGF");
    m.f_plus = Column(mtz, "F(+)");
    m.sigma_f_plus = Column(mtz, "SIGF(+)");
    m.f_minus = Column(mtz, "F(-)");
    m.sigma_f_minus = Column(mtz, "SIGF(-)");
    m.free_flag = FlagColumn(mtz, "FreeR_flag");
    return m;
}

MapCoefficientsMessage MapCoefficientsFromMtz(const std::string &path) {
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    MapCoefficientsMessage c;
    c.unit_cell = ToUnitCell(mtz.cell);
    c.space_group_name = SpaceGroupName(mtz);
    c.wavelength_A = static_cast<float>(mtz.datasets.back().wavelength);
    c.h = IntColumn(mtz, "H");
    c.k = IntColumn(mtz, "K");
    c.l = IntColumn(mtz, "L");
    c.fp = Column(mtz, "FP");
    c.fc = Column(mtz, "FC");
    c.phic = Column(mtz, "PHIC");
    c.fwt = Column(mtz, "FWT");
    c.phwt = Column(mtz, "PHWT");
    c.delfwt = Column(mtz, "DELFWT");
    c.phdelwt = Column(mtz, "PHDELWT");
    c.fom = Column(mtz, "FOM");
    c.free_flag = FlagColumn(mtz, "FREE");
    return c;
}

PlacedModelMessage PlacedModelFromStructure(const gemmi::Structure &st, const std::string &data_space_group_name) {
    PlacedModelMessage p;
    p.name = st.name;
    p.unit_cell = ToUnitCell(st.cell);
    p.space_group_name = data_space_group_name;
    p.model_space_group_name = st.spacegroup_hm;
    bool has_aniso = false;
    for (const gemmi::Model &m : st.models)
        for (const gemmi::Chain &ch : m.chains)
            for (const gemmi::Residue &r : ch.residues)
                for (const gemmi::Atom &a : r.atoms) {
                    p.model_number.push_back(m.num);
                    p.chain.push_back(ch.name);
                    p.residue_name.push_back(r.name);
                    p.seq_num.push_back(*r.seqid.num);
                    p.icode.push_back(static_cast<uint8_t>(r.seqid.icode));
                    p.het_flag.push_back(static_cast<uint8_t>(r.het_flag));
                    p.atom_name.push_back(a.name);
                    p.element.push_back(a.element.name());
                    p.altloc.push_back(static_cast<uint8_t>(a.altloc));
                    p.charge.push_back(a.charge);
                    p.x.push_back(a.pos.x);
                    p.y.push_back(a.pos.y);
                    p.z.push_back(a.pos.z);
                    p.b_iso.push_back(a.b_iso);
                    p.occupancy.push_back(a.occ);
                    p.aniso.insert(p.aniso.end(), {a.aniso.u11, a.aniso.u22, a.aniso.u33,
                                                   a.aniso.u12, a.aniso.u13, a.aniso.u23});
                    has_aniso = has_aniso || a.aniso.nonzero();
                }
    if (!has_aniso)
        p.aniso.clear();
    return p;
}

void FillRugnuxResultBlocks(const std::string &prefix, EndMessage &end) {
    if (std::filesystem::exists(prefix + ".mtz"))
        end.merged = MergedReflectionsFromMtz(prefix + ".mtz");
    if (std::filesystem::exists(prefix + "_maps.mtz"))
        end.map_coefficients = MapCoefficientsFromMtz(prefix + "_maps.mtz");
    if (std::filesystem::exists(prefix + "_model.cif")) {
        const gemmi::Structure st = gemmi::read_structure_gz(prefix + "_model.cif");
        end.placed_model = PlacedModelFromStructure(st, end.merged ? end.merged->space_group_name
                                                                   : st.spacegroup_hm);
    }
}