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

// Refinement (image_analysis/structure_refinement/Refinement.h) of a model against a rugnux merged MTZ, without
// processing anything: for measuring the refinement on finished runs (a battery run's p.mtz and p_model.cif)
// against another program on the same files. Prints the outcome, and one JSON line for scripts.

#include <cmath>
#include <cstdlib>
#include <iostream>
#include <sstream>

#include <gemmi/mmread_gz.hpp>
#include <gemmi/mtz.hpp>
#include <gemmi/to_cif.hpp>
#include <gemmi/to_mmcif.hpp>
#include <fstream>
#include <nlohmann/json.hpp>

#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"
#include "../image_analysis/structure_refinement/Refinement.h"

namespace {

void usage() {
    std::cout << "Usage: rugnux_refine_test [options] <merged.mtz> <model.cif|.pdb>\n"
                 "  -N <threads>        CPU threads (default 8)\n"
                 "  --no-gpu            the CPU engine\n"
                 "  --placement-only    rigid bodies and group B only\n"
                 "  --score-only        score the model as given, refine nothing\n"
                 "  --no-group-b        no B shift per residue in the placement stage\n"
                 "  --no-hydrogens      no riding hydrogens\n"
                 "  --ladder <a,b,...>  the first weight trials, times the gradient-balanced weight\n"
                 "  --adp-ladder <a,..> each with the B factors' weight that weight times these (0: B held)\n"
                 "  --steps <n>         the trajectory's steps at most\n"
                 "  --final-ladder <a,..> the final trials, times the trajectory's weights\n"
                 "  --max-weight <x>    the positions' weight at most, times the gradient-balanced weight\n"
                 "  --cycles <n>        macro-cycles per weight trial (default as in RefinementSettings)\n"
                 "  --xyz-iterations <n>, --b-iterations <n>  L-BFGS iterations per macro-cycle\n"
                 "  -o <file.cif>       write the refined model\n"
                 "  --dump <file>       with --score-only: h k l |Fcalc| phi |Fmask| phi k_total D var per reflection\n"
                 "  --trials-at-once <n> at most n weight trials side by side (1: one after another)\n"
                 "  -v                  debug lines (TIMING per refinement stage)\n";
}

std::vector<double> ParseList(const std::string &s) {
    std::vector<double> out;
    std::stringstream ss(s);
    std::string item;
    while (std::getline(ss, item, ','))
        out.push_back(std::stod(item));
    return out;
}

}  // namespace

int main(int argc, char **argv) {
    RefinementSettings settings;
    settings.nthreads = 8;
    std::string out_path, dump_path;
    bool verbose = false;
    std::vector<std::string> files;
    for (int i = 1; i < argc; i++) {
        const std::string a = argv[i];
        if (a == "-N" && i + 1 < argc)
            settings.nthreads = std::stoul(argv[++i]);
        else if (a == "--no-gpu") {
            settings.gpu = false;
            hide_gpus();
        } else if (a == "--no-group-b")
            settings.group_b = false;
        else if (a == "--no-hydrogens")
            settings.riding_hydrogens = false;
        else if (a == "--placement-only")
            settings.restrained = false;
        else if (a == "--score-only")
            settings.score_only = true;
        else if (a == "--cycles" && i + 1 < argc)
            settings.macro_cycles = std::stoi(argv[++i]);
        else if (a == "--xyz-iterations" && i + 1 < argc)
            settings.xyz_iterations = std::stoi(argv[++i]);
        else if (a == "--adp-ladder" && i + 1 < argc)
            settings.adp_ladder = ParseList(argv[++i]);
        else if (a == "--b-iterations" && i + 1 < argc)
            settings.b_iterations = std::stoi(argv[++i]);
        else if (a == "--ladder" && i + 1 < argc)
            settings.weight_ladder = ParseList(argv[++i]);
        else if (a == "--steps" && i + 1 < argc)
            settings.max_steps = std::stoi(argv[++i]);
        else if (a == "--final-ladder" && i + 1 < argc)
            settings.final_ladder = ParseList(argv[++i]);
        else if (a == "--max-weight" && i + 1 < argc)
            settings.max_weight = std::stod(argv[++i]);
        else if (a == "-v")
            verbose = true;
        else if (a == "--trials-at-once" && i + 1 < argc)
            settings.max_trials_at_once = std::stoul(argv[++i]);
        else if (a == "--dump" && i + 1 < argc)
            dump_path = argv[++i];
        else if (a == "-o" && i + 1 < argc)
            out_path = argv[++i];
        else if (!a.empty() && a[0] == '-') {
            usage();
            return EXIT_FAILURE;
        } else
            files.push_back(a);
    }
    if (files.size() != 2) {
        usage();
        return EXIT_FAILURE;
    }

    gemmi::Mtz mtz = gemmi::read_mtz_file(files[0]);
    const gemmi::SpaceGroup *sg = mtz.spacegroup;
    if (sg == nullptr) {
        std::cerr << "No space group in " << files[0] << std::endl;
        return EXIT_FAILURE;
    }
    const int ih = mtz.column_with_label("H")->idx, ik = mtz.column_with_label("K")->idx,
              il = mtz.column_with_label("L")->idx, iF = mtz.column_with_label("F")->idx,
              isF = mtz.column_with_label("SIGF")->idx, ifree = mtz.column_with_label("FreeR_flag")->idx;
    const gemmi::Mtz::Column *fp = mtz.column_with_label("F(+)"), *fm = mtz.column_with_label("F(-)");
    std::vector<MergedReflection> merged;
    for (size_t r = 0; r < static_cast<size_t>(mtz.nreflections); r++) {
        const float *row = &mtz.data[r * mtz.columns.size()];
        MergedReflection m;
        m.h = static_cast<int>(row[ih]);
        m.k = static_cast<int>(row[ik]);
        m.l = static_cast<int>(row[il]);
        m.F = row[iF];
        m.sigmaF = row[isF];
        m.rfree_flag = row[ifree] == 0;
        m.d = static_cast<float>(mtz.cell.calculate_d({{m.h, m.k, m.l}}));
        if (fp != nullptr && fm != nullptr) {
            m.F_plus = row[fp->idx];
            m.F_minus = row[fm->idx];
        }
        merged.push_back(m);
    }
    gemmi::Structure st = gemmi::read_structure_gz(files[1]);

    Logger logger("rugnux_refine_test");
    logger.Verbose(verbose);
    const RefinementResult r = RefineModel(st, merged, mtz.cell, *sg, 1.0, logger, settings);
    if (!r.ok) {
        std::cerr << "Refinement failed: " << r.failure_reason << std::endl;
        return EXIT_FAILURE;
    }
    if (!out_path.empty())
    {
        std::ofstream f(out_path);
        gemmi::cif::write_cif_to_stream(f, gemmi::make_mmcif_document(*r.model));
    }

    if (!dump_path.empty()) {
        std::ofstream f(dump_path);
        const RefinementScore &sc = r.score_start;
        for (size_t m = 0; m < sc.hkl.size(); m++)
            f << sc.hkl[m][0] << " " << sc.hkl[m][1] << " " << sc.hkl[m][2] << " " << std::abs(sc.fcalc[m]) << " "
              << std::arg(sc.fcalc[m]) << " " << std::abs(sc.fmask[m]) << " " << std::arg(sc.fmask[m]) << " "
              << sc.k_total[m] << " " << sc.d[m] << " " << sc.var[m] << "\n";
    }
    int kinds[5] = {0, 0, 0, 0, 0};
    for (const DensityBlob &b : r.blobs)
        kinds[static_cast<int>(b.kind)]++;
    nlohmann::json blobs = nlohmann::json::array();
    for (const DensityBlob &b : r.blobs)
        blobs.push_back({{"kind", DensityBlobKindName(b.kind)}, {"x", b.centroid.x}, {"y", b.centroid.y},
                         {"z", b.centroid.z}, {"peak", b.peak_sigma}, {"volume", b.volume}, {"score", b.score},
                         {"nearest", b.nearest_atom}, {"distance", b.nearest_atom_distance},
                         {"anomalous", b.anomalous_sigma}});
    nlohmann::json j = {{"d_min", r.d_min}, {"gpu", r.on_gpu}, {"seconds", r.seconds},
                        {"r_work_start", r.r_work_start}, {"r_free_start", r.r_free_start},
                        {"r_work_placement", r.r_work_placement}, {"r_free_placement", r.r_free_placement},
                        {"r_work", r.r_work}, {"r_free", r.r_free}, {"weight", r.weight}, {"adp_weight", r.adp_weight},
                        {"rmsd_bonds", r.geometry.rmsd_bonds}, {"rmsd_angles", r.geometry.rmsd_angles},
                        {"rmsd_bonds_start", r.geometry_start.rmsd_bonds},
                        {"rmsd_angles_start", r.geometry_start.rmsd_angles}, {"clashes", r.geometry.clashes},
                        {"atoms", r.atoms}, {"refined_atoms", r.refined_atoms}, {"fixed_residues", r.fixed_residues},
                        {"blobs", blobs}};
    nlohmann::json trials = nlohmann::json::array();
    for (const RefinementWeightTrial &t : r.trials)
        trials.push_back({{"weight", t.weight}, {"adp_weight", t.adp_weight}, {"r_work", t.r_work}, {"r_free", t.r_free},
                          {"rmsd_bonds", t.geometry.rmsd_bonds}, {"rmsd_angles", t.geometry.rmsd_angles}});
    j["trials"] = trials;
    const RefinementScore &sc = r.score_start;
    nlohmann::json shells = nlohmann::json::array();
    for (const RefinementShell &sh : sc.shells)
        shells.push_back({{"d_max", sh.d_max}, {"d_min", sh.d_min}, {"n_work", sh.n_work}, {"n_free", sh.n_free},
                          {"r_work", sh.r_work}, {"r_free", sh.r_free}});
    j["score_start"] = {{"k_overall", sc.k_overall}, {"b_iso", sc.b_iso}, {"b_aniso", sc.b_aniso}, {"k_sol", sc.k_sol},
                        {"b_sol", sc.b_sol}, {"mask_spacing", sc.mask_spacing}, {"r_work", sc.r_work},
                        {"r_free", sc.r_free}, {"n_work", sc.n_work}, {"n_free", sc.n_free}, {"nll_work", sc.nll_work},
                        {"nll_free", sc.nll_free}, {"hydrogens", sc.hydrogens}, {"fom", sc.fom}, {"mean_fom", r.mean_fom},
                        {"rmsz_bonds", r.geometry_start.rmsz_bonds}, {"rmsz_angles", r.geometry_start.rmsz_angles},
                        {"shells", shells}};
    std::cout << "JSON " << j.dump() << std::endl;
    return EXIT_SUCCESS;
}