Jungfraujoch/image_analysis/scale_merge/MtzWriter.cpp

// SPDX-FileCopyrightText: 2025 Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-only

#include "MtzWriter.h"

#include <cstring>
#include <cmath>
#include <vector>
#include <stdexcept>

// CCP4 C library — lives inside the CMtz namespace when included from C++
#include "../../ccp4c/ccp4/cmtzlib.h"
#include "../../ccp4c/ccp4/csymlib.h"
#include "../../ccp4c/ccp4/mtzdata.h"

MtzWriteOptions::MtzWriteOptions(const DiffractionExperiment& experiment,
                                 const std::optional<UnitCell>& cell_override) {
    // Unit cell: prefer explicit override, then experiment setting, else default
    if (cell_override.has_value()) {
        cell = cell_override.value();
    } else if (experiment.GetUnitCell().has_value()) {
        cell = experiment.GetUnitCell().value();
    }

    // Space group
    auto sg = experiment.GetGemmiSpaceGroup();
    if (sg) {
        spacegroup_number = sg->ccp4;
        spacegroup_name = sg->xhm();
    }

    // Wavelength
    wavelength = experiment.GetWavelength_A();

    // Sample name → crystal name
    auto sample = experiment.GetSampleName();
    if (!sample.empty())
        crystal_name = sample;

    // File prefix → project name
    auto prefix = experiment.GetFilePrefix();
    if (!prefix.empty())
        project_name = prefix;
}

namespace {

/// Helper: create a fresh in-memory MTZ structure, set up symmetry,
/// crystal, dataset, and the requested columns.
/// Returns {mtz, columns} or throws on failure.
struct MtzSetup {
    CMtz::MTZ* mtz = nullptr;
    std::vector<CMtz::MTZCOL*> cols;
};

MtzSetup CreateMtzSkeleton(const MtzWriteOptions& opts,
                            const char labels[][31],
                            const char types[][3],
                            int ncol) {
    using namespace CMtz;

    // Allocate empty MTZ (0 crystals initially)
    int nset_zero = 0;
    MTZ* mtz = MtzMalloc(0, &nset_zero);
    if (!mtz)
        throw std::runtime_error("MtzMalloc failed");

    mtz->refs_in_memory = 1;  // we build everything in memory, then write

    // Title
    ccp4_lwtitl(mtz, opts.title.c_str(), 0);

    // Symmetry: use CCP4 spacegroup lookup
    CCP4SPG* spg = ccp4spg_load_by_ccp4_num(opts.spacegroup_number);
    if (!spg) {
        // Fallback to P1
        spg = ccp4spg_load_by_ccp4_num(1);
    }

    if (spg) {
        const int nsym = spg->nsymop;
        const int nsymp = spg->nsymop_prim;
        float rsymx[192][4][4];
        std::memset(rsymx, 0, sizeof(rsymx));

        for (int i = 0; i < nsym && i < 192; ++i) {
            for (int r = 0; r < 3; ++r) {
                for (int c = 0; c < 3; ++c)
                    rsymx[i][r][c] = static_cast<float>(spg->symop[i].rot[r][c]);
                rsymx[i][r][3] = static_cast<float>(spg->symop[i].trn[r]);
            }
            rsymx[i][3][3] = 1.0f;
        }

        char ltypex[2] = {spg->symbol_old[0], '\0'};
        char spgrnx[MAXSPGNAMELENGTH + 1];
        std::strncpy(spgrnx, spg->symbol_xHM, MAXSPGNAMELENGTH);
        spgrnx[MAXSPGNAMELENGTH] = '\0';
        char pgnamx[MAXPGNAMELENGTH + 1];
        std::strncpy(pgnamx, spg->pgname, MAXPGNAMELENGTH);
        pgnamx[MAXPGNAMELENGTH] = '\0';

        ccp4_lwsymm(mtz, nsym, nsymp, rsymx,
                     ltypex, spg->spg_ccp4_num, spgrnx, pgnamx);

        ccp4spg_free(&spg);
    }

    // Crystal + dataset
    float cell[6] = {
        static_cast<float>(opts.cell.a),
        static_cast<float>(opts.cell.b),
        static_cast<float>(opts.cell.c),
        static_cast<float>(opts.cell.alpha),
        static_cast<float>(opts.cell.beta),
        static_cast<float>(opts.cell.gamma)
    };

    MTZXTAL* xtal = MtzAddXtal(mtz,
                                opts.crystal_name.c_str(),
                                opts.project_name.c_str(),
                                cell);
    if (!xtal) {
        MtzFree(mtz);
        throw std::runtime_error("MtzAddXtal failed");
    }

    MTZSET* set = MtzAddDataset(mtz, xtal,
                                 opts.dataset_name.c_str(),
                                 opts.wavelength);
    if (!set) {
        MtzFree(mtz);
        throw std::runtime_error("MtzAddDataset failed");
    }

    // Add columns
    MtzSetup result;
    result.mtz = mtz;
    result.cols.resize(ncol, nullptr);

    for (int i = 0; i < ncol; ++i) {
        result.cols[i] = MtzAddColumn(mtz, set, labels[i], types[i]);
        if (!result.cols[i]) {
            MtzFree(mtz);
            throw std::runtime_error(std::string("MtzAddColumn failed for ") + labels[i]);
        }
    }

    return result;
}

} // namespace

bool WriteMtzIntensities(const std::string& filename,
                         const std::vector<MergedReflection>& merged,
                         const MtzWriteOptions& opts) {
    using namespace CMtz;

    if (merged.empty())
        return false;

    constexpr int NCOL = 5;
    const char labels[NCOL][31] = {"H", "K", "L", "IMEAN", "SIGIMEAN"};
    const char types[NCOL][3]   = {"H", "H", "H", "J",     "Q"};

    MtzSetup setup;
    try {
        setup = CreateMtzSkeleton(opts, labels, types, NCOL);
    } catch (...) {
        return false;
    }

    // Write reflections (1-indexed)
    const int nref = static_cast<int>(merged.size());
    for (int i = 0; i < nref; ++i) {
        float adata[NCOL];
        adata[0] = static_cast<float>(merged[i].h);
        adata[1] = static_cast<float>(merged[i].k);
        adata[2] = static_cast<float>(merged[i].l);
        adata[3] = static_cast<float>(merged[i].I);
        adata[4] = std::isfinite(static_cast<float>(merged[i].sigma))
                       ? static_cast<float>(merged[i].sigma)
                       : 0.0f;

        ccp4_lwrefl(setup.mtz, adata, setup.cols.data(), NCOL, i + 1);
    }

    // Write to disk
    setup.mtz->nref = nref;
    int ok = MtzPut(setup.mtz, filename.c_str());
    MtzFree(setup.mtz);
    return ok != 0;
}

bool WriteMtzAmplitudes(const std::string& filename,
                        const std::vector<FrenchWilsonReflection>& fw,
                        const MtzWriteOptions& opts) {
    using namespace CMtz;

    if (fw.empty())
        return false;

    constexpr int NCOL = 7;
    const char labels[NCOL][31] = {"H", "K", "L", "FP", "SIGFP", "IMEAN", "SIGIMEAN"};
    const char types[NCOL][3]   = {"H", "H", "H", "F",  "Q",     "J",     "Q"};

    MtzSetup setup;
    try {
        setup = CreateMtzSkeleton(opts, labels, types, NCOL);
    } catch (...) {
        return false;
    }

    const int nref = static_cast<int>(fw.size());
    for (int i = 0; i < nref; ++i) {
        float adata[NCOL];
        adata[0] = static_cast<float>(fw[i].h);
        adata[1] = static_cast<float>(fw[i].k);
        adata[2] = static_cast<float>(fw[i].l);
        adata[3] = static_cast<float>(fw[i].F);
        adata[4] = static_cast<float>(fw[i].sigmaF);
        adata[5] = static_cast<float>(fw[i].I);
        adata[6] = std::isfinite(static_cast<float>(fw[i].sigmaI))
                       ? static_cast<float>(fw[i].sigmaI)
                       : 0.0f;

        ccp4_lwrefl(setup.mtz, adata, setup.cols.data(), NCOL, i + 1);
    }

    setup.mtz->nref = nref;
    int ok = MtzPut(setup.mtz, filename.c_str());
    MtzFree(setup.mtz);
    return ok != 0;
}