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

#include "WriteReflections.h"

#include <cmath>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <stdexcept>
#include <ctime>
#include <chrono>

#include <gemmi/mtz.hpp>

#include "../common/GitInfo.h"

namespace {

/// Current date in ISO-8601 (YYYY-MM-DD) for the _audit block.
std::string CurrentDateISO() {
    auto now = std::chrono::system_clock::now();
    auto t   = std::chrono::system_clock::to_time_t(now);
    std::tm tm{};
#ifdef _WIN32
    gmtime_s(&tm, &t);
#else
    gmtime_r(&t, &tm);
#endif
    char buf[32];
    std::strftime(buf, sizeof(buf), "%Y-%m-%d", &tm);
    return buf;
}

/// Format a double with given decimal places; returns "?" for non-finite.
std::string Fmt(double val, int decimals = 4) {
    if (!std::isfinite(val))
        return "?";
    std::ostringstream ss;
    ss << std::fixed << std::setprecision(decimals) << val;
    return ss.str();
}

/// Quote a CIF string value; returns "?" for empty.
std::string CifStr(const std::string& s) {
    if (s.empty())
        return "?";
    // If it contains spaces or special chars, single-quote it
    if (s.find(' ') != std::string::npos ||
        s.find('\'') != std::string::npos ||
        s.find('#') != std::string::npos)
        return "'" + s + "'";
    return s;
}

} // namespace


void WriteMmcifReflections(const std::vector<MergedReflection> &reflections,
                           const UnitCell &unitCell,
                           const DiffractionExperiment &experiment,
                           const std::string &filename) {

    std::ofstream out(filename);
    if (!out)
        throw std::runtime_error("WriteMmcifReflections: cannot open " + filename);

    out << std::fixed;

    // ---------- data block ----------
    out << "data_sample" << "\n";
    out << "#\n";

    // ---------- _audit ----------
    out << "_audit.revision_id     1\n";
    out << "_audit.creation_date   " << CurrentDateISO() << "\n";
    out << "_audit.update_record   'Initial release'\n";
    out << "#\n";

    // ---------- _software ----------
    out << "_software.name         'Jungfraujoch'\n";

    out << "_software.version      " << CifStr(jfjoch_version()) << "\n";
    out << "_software.classification reduction\n";
    out << "#\n";

    // ---------- _cell ----------
    out << "_cell.length_a     " << Fmt(unitCell.a, 3) << "\n";
    out << "_cell.length_b     " << Fmt(unitCell.b, 3) << "\n";
    out << "_cell.length_c     " << Fmt(unitCell.c, 3) << "\n";
    out << "_cell.angle_alpha  " << Fmt(unitCell.alpha, 2) << "\n";
    out << "_cell.angle_beta   " << Fmt(unitCell.beta, 2) << "\n";
    out << "_cell.angle_gamma  " << Fmt(unitCell.gamma, 2) << "\n";

    auto *sg = gemmi::find_spacegroup_by_number(experiment.GetSpaceGroupNumber().value_or(1));
    if (sg == nullptr)
        throw std::runtime_error("WriteMmcifReflections: invalid space group number");

    // ---------- _symmetry ----------
    out << "_symmetry.space_group_name_H-M  " << CifStr(sg->hm) << "\n";
    out << "_symmetry.Int_Tables_number      " << sg->number << "\n";
    out << "#\n";

    // ---------- _diffrn_source / _diffrn_detector ----------
    if (!experiment.GetSourceName().empty())
        out << "_diffrn_source.pdbx_synchrotron_site  " << CifStr(experiment.GetSourceName()) << "\n";

    if (!experiment.GetInstrumentName().empty())
        out << "_diffrn_source.pdbx_synchrotron_beamline  " << CifStr(experiment.GetInstrumentName()) << "\n";

    out << "_diffrn_radiation_wavelength.wavelength  " << Fmt(experiment.GetWavelength_A(), 5) << "\n";
    out << "_diffrn_detector.detector  " << CifStr(experiment.GetDetectorDescription()) << "\n";

    // ---------- _refln loop ----------
    out << "loop_\n";
    out << "_refln.index_h\n";
    out << "_refln.index_k\n";
    out << "_refln.index_l\n";
    out << "_refln.intensity_meas\n";
    out << "_refln.intensity_sigma\n";
    out << "_refln.status_free\n";
    out << "_refln.status\n";

    for (const auto& r : reflections) {
        out << std::setw(5) << r.h << " "
            << std::setw(5) << r.k << " "
            << std::setw(5) << r.l << " "
            << std::setw(14) << Fmt(r.I, 4) << " "
            << std::setw(14) << Fmt(r.sigma, 4) << " "
            << (r.rfree_flag ? 1 : 0) << " "
            << "o"  // 'o' = observed
            << "\n";
    }

    out << "#\n";
    out << "# End of reflections\n";
    out.close();
}

void WriteMtzReflections(const std::vector<MergedReflection> &reflections,
                         const UnitCell &unitCell,
                         const DiffractionExperiment &experiment,
                         const std::string &filename) {
    gemmi::Mtz mtz;

    // Optional but recommended metadata
    mtz.spacegroup = gemmi::find_spacegroup_by_number(
        experiment.GetSpaceGroupNumber().value_or(1));
    mtz.set_cell_for_all(unitCell);

    // Add dataset
    gemmi::Mtz::Dataset& ds = mtz.add_dataset("native");
    ds.crystal_name = experiment.GetSampleName();
    ds.wavelength = experiment.GetWavelength_A();

    int dataset_id = ds.id;

    // Add columns
    mtz.add_column("H", 'H', dataset_id, -1, false);
    mtz.add_column("K", 'H', dataset_id, -1, false);
    mtz.add_column("L", 'H', dataset_id, -1, false);
    mtz.add_column("IMEAN", 'J', dataset_id, -1, false);
    mtz.add_column("SIGIMEAN", 'Q', dataset_id, -1, false);
    mtz.add_column("FreeR_flag", 'I', dataset_id, -1, false);

    // Number of rows
    mtz.nreflections = reflections.size();

    // Allocate data table
    mtz.data.reserve(reflections.size() * 6);

    for (const auto& r : reflections) {
        mtz.data.push_back(static_cast<float>(r.h));
        mtz.data.push_back(static_cast<float>(r.k));
        mtz.data.push_back(static_cast<float>(r.l));
        mtz.data.push_back(r.I);
        mtz.data.push_back(r.sigma);
        mtz.data.push_back(r.rfree_flag ? 1 : 0);
    }

    // Write MTZ
    mtz.write_to_file(filename);
}

void WriteHKLReflections(const std::vector<MergedReflection> &reflections,
                         const std::string &filename) {
    std::ofstream hkl_file(filename);
    if (!hkl_file)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "Cannot open file " + filename + " for writing");

    for (const auto &r: reflections)
        hkl_file << r.h << " " << r.k << " " << r.l << " "
                << r.I << " " << r.sigma << " "
                << (r.rfree_flag ? 1 : 0) << std::endl;

    hkl_file.close();
}

void WriteReflections(const std::vector<MergedReflection> &reflections,
                      const UnitCell &unitCell,
                      const DiffractionExperiment &experiment,
                      const std::string &filename) {
    switch (experiment.GetScalingSettings().GetFileFormat()) {
        case IntensityFormat::Text:
            WriteHKLReflections(reflections, filename + ".hkl");
            break;
        case IntensityFormat::mmCIF:
            WriteMmcifReflections(reflections, unitCell, experiment, filename + ".cif");
            break;
        case IntensityFormat::MTZ:
            WriteMtzReflections(reflections, unitCell, experiment, filename + ".mtz");
            break;
    }
}