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

#include "WriteRugnuxFiles.h"

#include <chrono>
#include <filesystem>

#include <gemmi/mtz.hpp>

#include "WriteModel.h"
#include "../common/Logger.h"

void WriteMergedMtz(const MergedReflectionsMessage &merged, const std::string &filename) {
    gemmi::Mtz mtz;

    // Optional but recommended metadata
    mtz.spacegroup = gemmi::find_spacegroup_by_name(merged.space_group_name);
    mtz.set_cell_for_all(merged.unit_cell);

    // Producing-software provenance in the MTZ header (title + HISTORY, the CCP4 convention).
    mtz.title = "Rugnux merged reflections";
    mtz.history = merged.history;

    // Add dataset. HKL_base first, so the data dataset is id 1: a dataset written as id 0 is the
    // reserved base dataset as far as CCP4's mtzlib is concerned, and the wavelength on it is
    // dropped - mtzinfo then reports the Cu K-alpha default 1.54187 A, which is what a program that
    // takes f'/f'' from the file would phase on.
    mtz.add_base();                 // the HKL_base dataset and the H K L columns
    gemmi::Mtz::Dataset& ds = mtz.add_dataset("native");
    ds.crystal_name = merged.crystal_name;
    ds.wavelength = merged.wavelength_A;

    const int dataset_id = ds.id;

    // One row per reflection in the CCP4 anomalous layout (IMEAN + I(+)/I(-), and the same split for
    // the French-Wilson amplitude), which aimless / ctruncate / mtz2sca / ANODE read directly.
    const bool has_anom = !merged.i_plus.empty();

    mtz.add_column("IMEAN", 'J', dataset_id, -1, false);
    mtz.add_column("SIGIMEAN", 'Q', dataset_id, -1, false);
    if (has_anom) {
        mtz.add_column("I(+)", 'K', dataset_id, -1, false);
        mtz.add_column("SIGI(+)", 'M', dataset_id, -1, false);
        mtz.add_column("I(-)", 'K', dataset_id, -1, false);
        mtz.add_column("SIGI(-)", 'M', dataset_id, -1, false);
    }
    mtz.add_column("F", 'F', dataset_id, -1, false);       // French-Wilson amplitude
    mtz.add_column("SIGF", 'Q', dataset_id, -1, false);
    if (has_anom) {
        mtz.add_column("F(+)", 'G', dataset_id, -1, false);
        mtz.add_column("SIGF(+)", 'L', dataset_id, -1, false);
        mtz.add_column("F(-)", 'G', dataset_id, -1, false);
        mtz.add_column("SIGF(-)", 'L', dataset_id, -1, false);
    }
    mtz.add_column("FreeR_flag", 'I', dataset_id, -1, false);

    const size_t n = merged.h.size();
    mtz.nreflections = static_cast<int>(n);
    mtz.data.reserve(n * 16);
    for (size_t i = 0; i < n; i++) {
        mtz.data.push_back(static_cast<float>(merged.h[i]));
        mtz.data.push_back(static_cast<float>(merged.k[i]));
        mtz.data.push_back(static_cast<float>(merged.l[i]));
        mtz.data.push_back(merged.i_mean[i]);
        mtz.data.push_back(merged.sigma_i_mean[i]);
        if (has_anom) {
            mtz.data.push_back(merged.i_plus[i]);
            mtz.data.push_back(merged.sigma_i_plus[i]);
            mtz.data.push_back(merged.i_minus[i]);
            mtz.data.push_back(merged.sigma_i_minus[i]);
        }
        mtz.data.push_back(merged.f[i]);
        mtz.data.push_back(merged.sigma_f[i]);
        if (has_anom) {
            mtz.data.push_back(merged.f_plus[i]);
            mtz.data.push_back(merged.sigma_f_plus[i]);
            mtz.data.push_back(merged.f_minus[i]);
            mtz.data.push_back(merged.sigma_f_minus[i]);
        }
        mtz.data.push_back(static_cast<float>(merged.free_flag[i]));
    }
    mtz.write_to_file(filename);
}

void WriteMapCoefficientsMtz(const MapCoefficientsMessage &coefficients, const std::string &filename) {
    gemmi::Mtz mtz(true);
    mtz.spacegroup = gemmi::find_spacegroup_by_name(coefficients.space_group_name);
    mtz.set_cell_for_all(coefficients.unit_cell);
    mtz.add_dataset("model_validation");
    mtz.datasets.back().wavelength = coefficients.wavelength_A;
    mtz.add_column("FP", 'F', -1, -1, false);
    mtz.add_column("FC", 'F', -1, -1, false);
    mtz.add_column("PHIC", 'P', -1, -1, false);
    mtz.add_column("FWT", 'F', -1, -1, false);
    mtz.add_column("PHWT", 'P', -1, -1, false);
    mtz.add_column("DELFWT", 'F', -1, -1, false);
    mtz.add_column("PHDELWT", 'P', -1, -1, false);
    // The figure of merit the coefficients carry, so the weighting can be read off - and undone -
    // from the file rather than having to be taken on trust.
    mtz.add_column("FOM", 'W', -1, -1, false);
    mtz.add_column("FREE", 'I', -1, -1, false);
    const size_t n = coefficients.h.size();
    mtz.data.reserve(n * 12);
    for (size_t i = 0; i < n; i++)
        mtz.data.insert(mtz.data.end(), {static_cast<float>(coefficients.h[i]),
                                         static_cast<float>(coefficients.k[i]),
                                         static_cast<float>(coefficients.l[i]),
                                         coefficients.fp[i], coefficients.fc[i], coefficients.phic[i],
                                         coefficients.fwt[i], coefficients.phwt[i],
                                         coefficients.delfwt[i], coefficients.phdelwt[i],
                                         coefficients.fom[i],
                                         static_cast<float>(coefficients.free_flag[i])});
    mtz.nreflections = static_cast<int>(n);
    mtz.write_to_file(filename);
}

gemmi::Structure PlacedModelStructure(const PlacedModelMessage &model) {
    gemmi::Structure st;
    st.name = model.name;
    st.cell = model.unit_cell;
    st.spacegroup_hm = model.model_space_group_name;
    const bool has_aniso = !model.aniso.empty();
    gemmi::Model *m = nullptr;
    gemmi::Chain *ch = nullptr;
    gemmi::Residue *res = nullptr;
    for (size_t i = 0; i < model.atom_name.size(); i++) {
        if (m == nullptr || m->num != model.model_number[i]) {
            m = &st.models.emplace_back(model.model_number[i]);
            ch = nullptr;
        }
        if (ch == nullptr || ch->name != model.chain[i]) {
            ch = &m->chains.emplace_back(model.chain[i]);
            res = nullptr;
        }
        const gemmi::SeqId seqid(model.seq_num[i], static_cast<char>(model.icode[i]));
        if (res == nullptr || res->name != model.residue_name[i] || res->seqid != seqid
            || res->het_flag != static_cast<char>(model.het_flag[i])) {
            res = &ch->residues.emplace_back(gemmi::ResidueId{seqid, {}, model.residue_name[i]});
            res->het_flag = static_cast<char>(model.het_flag[i]);
        }
        gemmi::Atom &a = res->atoms.emplace_back();
        a.name = model.atom_name[i];
        a.element = gemmi::Element(model.element[i]);
        a.altloc = static_cast<char>(model.altloc[i]);
        a.charge = static_cast<signed char>(model.charge[i]);
        a.pos = gemmi::Position(model.x[i], model.y[i], model.z[i]);
        a.b_iso = model.b_iso[i];
        a.occ = model.occupancy[i];
        if (has_aniso)
            a.aniso = {model.aniso[6 * i], model.aniso[6 * i + 1], model.aniso[6 * i + 2],
                       model.aniso[6 * i + 3], model.aniso[6 * i + 4], model.aniso[6 * i + 5]};
    }
    st.setup_cell_images();
    return st;
}

namespace {

double WallClockSeconds() {
    return std::chrono::duration<double>(std::chrono::system_clock::now().time_since_epoch()).count();
}

// Rename tmp onto final, unless final is there already and may not be overwritten.
void RenameIntoPlace(const std::string &tmp, const std::string &final_name, bool overwrite, Logger &logger) {
    if (!overwrite && std::filesystem::exists(final_name)) {
        std::filesystem::remove(tmp);
        logger.Error("{} already exists and overwrite is off - not written", final_name);
        return;
    }
    std::filesystem::rename(tmp, final_name);
    logger.Info("{} written, renamed into place at wall clock {:.3f} s", final_name, WallClockSeconds());
}

}  // namespace

void WriteRugnuxResultFiles(const EndMessage &end, const std::string &file_prefix, bool overwrite,
                            Logger &logger) {
    // A file that cannot be written costs only itself: the master is already in place.
    try {
        if (end.merged) {
            WriteMergedMtz(*end.merged, file_prefix + ".mtz.tmp");
            RenameIntoPlace(file_prefix + ".mtz.tmp", file_prefix + ".mtz", overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}.mtz: {}", file_prefix, e.what());
    }
    try {
        if (end.map_coefficients) {
            WriteMapCoefficientsMtz(*end.map_coefficients, file_prefix + "_maps.mtz.tmp");
            RenameIntoPlace(file_prefix + "_maps.mtz.tmp", file_prefix + "_maps.mtz", overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_maps.mtz: {}", file_prefix, e.what());
    }
    try {
        if (end.placed_model) {
            const PlacedModelMessage &model = *end.placed_model;
            // WritePlacedModel appends _model.cif and _model.pdb to the prefix it is given.
            const std::string tmp_prefix = file_prefix + ".tmp";
            WritePlacedModel(PlacedModelStructure(model), model.unit_cell,
                             *gemmi::find_spacegroup_by_name(model.space_group_name), tmp_prefix, logger);
            for (const std::string suffix : {"_model.cif", "_model.pdb"})
                if (std::filesystem::exists(tmp_prefix + suffix))
                    RenameIntoPlace(tmp_prefix + suffix, file_prefix + suffix, overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_model.cif: {}", file_prefix, e.what());
    }
}