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

#include "WriteModel.h"

#include <algorithm>
#include <fstream>
#include <sstream>

#include <gemmi/to_mmcif.hpp>  // make_mmcif_document
#include <gemmi/to_cif.hpp>    // write_cif_to_stream
#include <gemmi/to_pdb.hpp>    // write_pdb
#include <gemmi/polyheur.hpp>  // setup_entities

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

namespace {

// Do the coordinates of a model built in `coordinates_group` obey `label` as well? They do when every
// operation of `label` is one of theirs, i.e. when `label` is a subgroup: the file then holds more
// than one asymmetric unit of it, which is allowed and is what a model always does in P1. They do not
// when `label` has an operation of its own, and then the label is a lie a reader acts on - it
// generates atoms from that operation, and the crystal it builds is not the one in the file.
bool coordinates_obey(const gemmi::SpaceGroup &coordinates_group, const gemmi::SpaceGroup &label) {
    const std::vector<gemmi::Op> coordinate_ops = coordinates_group.operations().all_ops_sorted();
    const std::vector<gemmi::Op> label_ops = label.operations().all_ops_sorted();
    return std::includes(coordinate_ops.begin(), coordinate_ops.end(),
                         label_ops.begin(), label_ops.end());
}

// `placed` in `cell`, labelled with `space_group` wherever its coordinates obey it (see WritePlacedModel).
// takes_model_group: the model's own group was kept instead.
gemmi::Structure LabelledModel(const gemmi::Structure &placed, const UnitCell &cell,
                               const gemmi::SpaceGroup &space_group, bool &takes_model_group) {
    gemmi::Structure st = placed;
    // The cell is the data's, taken from the same value the reflection files are written from: the
    // coordinates were re-fractionalized into it and genuinely sit there.
    //
    // The group is the data's too wherever the coordinates obey it, which is the ordinary case and the
    // one that matters - the label then matches the .mtz beside it, the enantiomorph in particular,
    // since --model can adopt the model's and that is neither the data's original label nor,
    // necessarily, the one the input model arrived with. Where they do not obey it the model's own
    // group is written instead. Data merged in a supergroup of the model's - a run that over-merged
    // across a pseudo-symmetry twofold, or a hand the model's fit did not earn - would otherwise put
    // operations on the file that its contents contradict, and a refinement program expands them: the
    // structure it refines is not the model, and it says so nowhere. A pair whose two files disagree
    // is worth a warning; a coordinate file that is quietly wrong on its own is not worth writing.
    const gemmi::SpaceGroup *model_group = st.find_spacegroup();
    takes_model_group = model_group != nullptr && !coordinates_obey(*model_group, space_group);
    st.cell = cell;
    st.spacegroup_hm = (takes_model_group ? *model_group : space_group).xhm();
    st.setup_cell_images();
    return st;
}

// The TLS groups of the model as REMARK 3 TLS records (the layout REFMAC writes and gemmi reads back), in place of
// any REMARK 3 the model was read with: gemmi's PDB writer writes none of them. Nothing where it has no groups.
std::string WithTlsRemarks(const gemmi::Structure &st, const std::string &pdb) {
    const std::vector<gemmi::TlsGroup> *groups = st.meta.get_tls_groups();
    if (groups == nullptr || groups->empty())
        return pdb;
    std::ostringstream r;
    auto line = [&r](const std::string &text) { r << fmt::format("{:<80}\n", "REMARK   3" + text); };
    line("");
    line(" REFINEMENT.");
    line("   PROGRAM     : rugnux");
    line("");
    line("  DATA USED IN REFINEMENT.");
    line("");
    line("  TLS DETAILS");
    line(fmt::format("   NUMBER OF TLS GROUPS  : {:4}", groups->size()));
    for (const gemmi::TlsGroup &g : *groups) {
        line("");
        line(fmt::format("   TLS GROUP : {:5}", g.id));
        line(fmt::format("    NUMBER OF COMPONENTS GROUP : {:4}", g.selections.size()));
        line("    COMPONENTS        C SSSEQI   TO  C SSSEQI");
        for (const gemmi::TlsGroup::Selection &sel : g.selections)
            line(fmt::format("    RESIDUE RANGE :{:>5}{:>6}    {:>5}{:>6}", sel.chain, sel.res_begin.str(), sel.chain,
                             sel.res_end.str()));
        line(fmt::format("    ORIGIN FOR THE GROUP (A):{:9.4f}{:9.4f}{:9.4f}", g.origin.x, g.origin.y, g.origin.z));
        line("    T TENSOR");
        line(fmt::format("      T11: {:8.4f} T22: {:8.4f}", g.T.u11, g.T.u22));
        line(fmt::format("      T33: {:8.4f} T12: {:8.4f}", g.T.u33, g.T.u12));
        line(fmt::format("      T13: {:8.4f} T23: {:8.4f}", g.T.u13, g.T.u23));
        line("    L TENSOR");
        line(fmt::format("      L11: {:8.4f} L22: {:8.4f}", g.L.u11, g.L.u22));
        line(fmt::format("      L33: {:8.4f} L12: {:8.4f}", g.L.u33, g.L.u12));
        line(fmt::format("      L13: {:8.4f} L23: {:8.4f}", g.L.u13, g.L.u23));
        line("    S TENSOR");
        for (int i = 0; i < 3; i++)
            line(fmt::format("      S{}1: {:8.4f} S{}2: {:8.4f} S{}3: {:8.4f}", i + 1, g.S[i][0], i + 1, g.S[i][1], i + 1,
                             g.S[i][2]));
    }
    // before the crystal's record; the model's own REMARK 3 left out
    std::istringstream in(pdb);
    std::ostringstream out;
    std::string text;
    bool written = false;
    while (std::getline(in, text)) {
        if (text.rfind("REMARK   3", 0) == 0)
            continue;
        if (!written && (text.rfind("CRYST1", 0) == 0 || text.rfind("ATOM", 0) == 0 || text.rfind("HETATM", 0) == 0)) {
            out << r.str();
            written = true;
        }
        out << text << '\n';
    }
    return out.str();
}

}  // namespace

void WritePlacedModel(const gemmi::Structure &placed,
                      const UnitCell &cell,
                      const gemmi::SpaceGroup &space_group,
                      const std::string &output_prefix,
                      Logger &logger) {
    const std::string cif_path = output_prefix + "_model.cif";
    // PDB as well as mmCIF, because the fragment-screening tools this file exists for take a PDB:
    // PanDDA's per-dataset input is <name>.pdb beside <name>.mtz, and dimple produces the same pair.
    // Same coordinates, same frame, both formats - a PDB cannot hold every cell, so it is written
    // when it can be and skipped, with a line saying so, when it cannot.
    const std::string pdb_path = output_prefix + "_model.pdb";

    bool takes_model_group = false;
    gemmi::Structure st = LabelledModel(placed, cell, space_group, takes_model_group);
    const gemmi::SpaceGroup &label = takes_model_group ? *st.find_spacegroup() : space_group;
    // Fills in entity types and label_asym_id for a model read from a PDB, which carries neither.
    // Both are no-ops where the input already had them, i.e. for an mmCIF input.
    // setup_entities and make_mmcif_document can both throw, and this runs BEFORE the reflection files
    // are written: a convenience deliverable must not be able to take the run's actual output with it.
    // The map writer does the same.
    try {
        gemmi::setup_entities(st);
        std::ofstream os(cif_path);
        gemmi::cif::write_cif_to_stream(os, gemmi::make_mmcif_document(st));
        if (!os) {
            logger.Error("Model validation: cannot write the placed model to {}", cif_path);
            return;
        }
    } catch (const std::exception &e) {
        logger.Warning("Model validation: could not write the placed model to {} ({})",
                       cif_path, e.what());
        return;
    }
    bool pdb_written = false;
    try {
        std::ostringstream pdb;
        gemmi::write_pdb(st, pdb);
        std::ofstream os(pdb_path);
        os << WithTlsRemarks(st, pdb.str());
        pdb_written = static_cast<bool>(os);
    } catch (const std::exception &e) {
        logger.Warning("Model validation: the placed model could not also be written as PDB ({}); "
                       "{} has it", e.what(), cif_path);
    }
    if (!takes_model_group) {
        logger.Info("Model validation: the model as placed against these data written to {}{} "
                    "(cell and space group {} as in the reflection files)",
                    cif_path, pdb_written ? " and " + pdb_path : "", space_group.short_name());
    } else {
        logger.Warning("Model validation: the model as placed written to {}{} in ITS OWN space group "
                       "{}, not the {} the reflection files carry - its coordinates do not obey {}, so "
                       "the two files describe the same crystal in different symmetry and cannot be "
                       "refined against each other as they stand",
                       cif_path, pdb_written ? " and " + pdb_path : "",
                       label.short_name(), space_group.short_name(), space_group.short_name());
    }
}

std::string PlacedModelMmcif(const gemmi::Structure &placed, const UnitCell &cell,
                             const gemmi::SpaceGroup &space_group) {
    bool takes_model_group = false;
    gemmi::Structure st = LabelledModel(placed, cell, space_group, takes_model_group);
    gemmi::setup_entities(st);
    std::ostringstream os;
    gemmi::cif::write_cif_to_stream(os, gemmi::make_mmcif_document(st));
    return os.str();
}