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

#include "RugnuxReflectionFile.h"

#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iterator>
#include <stdexcept>

#include <gemmi/cif2mtz.hpp>
#include <gemmi/read_cif.hpp>
#include <gemmi/refln.hpp>

#include "RugnuxModelFile.h"   // GunzipIfCompressed

namespace {

std::string Lower(std::string s) {
    std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
    return s;
}

// The sigma of column `c`: the type-Q column named SIG<label>, else the type-Q column right after it.
const gemmi::Mtz::Column *SigmaOf(const gemmi::Mtz &mtz, const gemmi::Mtz::Column &c, char type) {
    if (const gemmi::Mtz::Column *s = mtz.column_with_label("SIG" + c.label, nullptr, type))
        return s;
    if (c.idx + 1 < mtz.columns.size() && mtz.columns[c.idx + 1].type == type)
        return &mtz.columns[c.idx + 1];
    return nullptr;
}

// Labels of amplitudes that are not observations: calculated structure factors and map coefficients.
bool CalculatedAmplitude(const std::string &label) {
    const std::string l = Lower(label);
    for (const char *p : {"fc", "fwt", "delfwt", "2fofc", "fofc", "f-model", "fmodel", "fcalc", "fpart", "fmask",
                          "fom", "fdiff", "fano", "dano", "fem"})
        if (l.starts_with(p))
            return true;
    return false;
}

// A value column of `type` with its sigma: the names in `preferred` first, then any (that `accept` takes).
template <typename Accept>
std::pair<const gemmi::Mtz::Column *, const gemmi::Mtz::Column *>
FindPair(const gemmi::Mtz &mtz, char type, std::initializer_list<const char *> preferred, Accept accept) {
    for (const char *label : preferred)
        if (const gemmi::Mtz::Column *c = mtz.column_with_label(label, nullptr, type))
            if (const gemmi::Mtz::Column *s = SigmaOf(mtz, *c, 'Q'))
                return {c, s};
    for (const gemmi::Mtz::Column &c : mtz.columns)
        if (c.type == type && accept(c.label))
            if (const gemmi::Mtz::Column *s = SigmaOf(mtz, c, 'Q'))
                return {&c, s};
    return {nullptr, nullptr};
}

// The free-set flag column: FreeR_flag and the other usual names first, then any type-I column with "free" in its name.
const gemmi::Mtz::Column *FindFree(const gemmi::Mtz &mtz) {
    for (const char *label : {"FreeR_flag", "FREE", "FREER", "RFREE", "R-free-flags", "FreeRflag", "FREE_FLAG"})
        if (const gemmi::Mtz::Column *c = mtz.column_with_label(label, nullptr, 'I'))
            return c;
    for (const gemmi::Mtz::Column &c : mtz.columns)
        if (c.type == 'I' && Lower(c.label).find("free") != std::string::npos)
            return &c;
    return nullptr;
}

// The anomalous pair of `value_type` (G or K) with sigmas of `sigma_type` (L or M), named X(+) and X(-).
bool FindAnomalous(const gemmi::Mtz &mtz, char value_type, char sigma_type, const gemmi::Mtz::Column *out[4]) {
    for (const gemmi::Mtz::Column &c : mtz.columns) {
        if (c.type != value_type || !c.label.ends_with("(+)"))
            continue;
        const std::string base = c.label.substr(0, c.label.size() - 3);
        const gemmi::Mtz::Column *m = mtz.column_with_label(base + "(-)", nullptr, value_type);
        const gemmi::Mtz::Column *sp = SigmaOf(mtz, c, sigma_type);
        const gemmi::Mtz::Column *sm = m != nullptr ? SigmaOf(mtz, *m, sigma_type) : nullptr;
        if (m != nullptr && sp != nullptr && sm != nullptr) {
            out[0] = &c;
            out[1] = sp;
            out[2] = m;
            out[3] = sm;
            return true;
        }
    }
    return false;
}

// gemmi's SF-mmCIF -> MTZ conversion with _refln.status preferred to _refln.pdbx_r_free_flag for the free set
// where status marks any reflection free ('f'), as LoadFCalcFromMtz.cpp reads a deposition.
std::vector<std::string> SfCifSpec(const gemmi::cif::Loop &loop) {
    bool status_has_free = false;
    const int status = loop.find_tag("_refln.status");
    if (status >= 0)
        for (size_t i = status; i < loop.values.size(); i += loop.width())
            if (gemmi::cif::as_string(loop.values[i]) == "f") {
                status_has_free = true;
                break;
            }
    std::vector<std::string> spec;
    for (const char **line = gemmi::CifToMtz::default_spec(true); *line != nullptr; ++line) {
        const std::string l = *line;
        if ((status_has_free && l.starts_with("pdbx_r_free_flag ")) || (!status_has_free && l.starts_with("status ")))
            continue;
        spec.push_back(l);
    }
    return spec;
}

}  // namespace

RugnuxReflections RugnuxReflectionsFromMtz(const gemmi::Mtz &mtz, const std::string &name) {
    RugnuxReflections out;
    out.name = name;
    out.cell = mtz.get_cell();
    if (mtz.spacegroup == nullptr)
        throw std::runtime_error(name + " names no space group");
    out.space_group = mtz.spacegroup->xhm();
    for (const gemmi::Mtz::Dataset &ds : mtz.datasets)
        if (ds.wavelength > 0)
            out.wavelength_A = ds.wavelength;

    auto [value, sigma] = FindPair(mtz, 'F', {"F", "FP", "FOBS", "F-obs", "FMEAN", "F_meas", "Fobs"},
                                   [](const std::string &l) { return !CalculatedAmplitude(l); });
    if (value == nullptr) {
        std::tie(value, sigma) = FindPair(mtz, 'J', {"IMEAN", "I", "IOBS", "I-obs", "Iobs", "I_meas"},
                                          [](const std::string &) { return true; });
        out.amplitudes = false;
    }
    const gemmi::Mtz::Column *anom[4] = {nullptr, nullptr, nullptr, nullptr};
    if (value == nullptr) {
        out.amplitudes = FindAnomalous(mtz, 'G', 'L', anom);
        if (!out.amplitudes && !FindAnomalous(mtz, 'K', 'M', anom))
            throw std::runtime_error(name + " holds no amplitudes or intensities with their sigmas");
        out.value_column = anom[0]->label + "," + anom[2]->label;
        out.sigma_column = anom[1]->label + "," + anom[3]->label;
    } else {
        out.value_column = value->label;
        out.sigma_column = sigma->label;
    }

    const gemmi::Mtz::Column *free = FindFree(mtz);
    std::vector<int> raw_free;
    const size_t stride = mtz.columns.size();
    for (size_t r = 0; r < static_cast<size_t>(mtz.nreflections); r++) {
        const float *row = &mtz.data[r * stride];
        float v = NAN, s = NAN;
        if (value != nullptr) {
            v = row[value->idx];
            s = row[sigma->idx];
        } else {
            // the mean of the two hands where both were measured, the one measured otherwise
            const float vp = row[anom[0]->idx], sp = row[anom[1]->idx], vm = row[anom[2]->idx], sm = row[anom[3]->idx];
            if (std::isfinite(vp) && std::isfinite(vm)) {
                v = 0.5f * (vp + vm);
                s = 0.5f * std::sqrt(sp * sp + sm * sm);
            } else if (std::isfinite(vp)) {
                v = vp;
                s = sp;
            } else {
                v = vm;
                s = sm;
            }
        }
        if (!std::isfinite(v))
            continue;
        MergedReflection m;
        m.h = static_cast<int32_t>(row[0]);
        m.k = static_cast<int32_t>(row[1]);
        m.l = static_cast<int32_t>(row[2]);
        m.d = static_cast<float>(out.cell.calculate_d({{m.h, m.k, m.l}}));
        if (out.amplitudes) {
            m.F = v;
            m.sigmaF = s;
        } else {
            m.I = v;
            m.sigma = s;
        }
        out.reflections.push_back(m);
        if (free != nullptr) {
            const float f = row[free->idx];
            raw_free.push_back(std::isfinite(f) ? static_cast<int>(std::lround(f)) : -1);
        }
    }
    if (out.reflections.empty())
        throw std::runtime_error(name + " holds no reflection with a value in " + out.value_column);

    if (free != nullptr) {
        size_t zeros = 0, flagged = 0;
        for (int f : raw_free) {
            zeros += f == 0;
            flagged += f >= 0;
        }
        const bool free_is_zero = zeros * 2 <= flagged;
        for (size_t i = 0; i < out.reflections.size(); i++) {
            const bool is_free = raw_free[i] >= 0 && (raw_free[i] == 0) == free_is_zero;
            out.reflections[i].rfree_flag = is_free;
            out.n_free += is_free;
        }
        // a column that marks nothing free is no free set
        if (out.n_free > 0)
            out.free_column = free->label;
    }
    return out;
}

RugnuxReflections ParseRugnuxReflections(const std::string &bytes, const std::string &name) {
    const std::string text = GunzipIfCompressed(bytes);
    if (text.size() >= 4 && std::memcmp(text.data(), "MTZ ", 4) == 0) {
        gemmi::Mtz mtz;
        gemmi::MemoryStream stream(text.data(), text.size());
        mtz.read_stream(stream, true);
        RugnuxReflections out = RugnuxReflectionsFromMtz(mtz, name);
        out.source = "MTZ";
        return out;
    }
    gemmi::cif::Document doc = gemmi::read_cif_from_memory(text.data(), text.size(), name.c_str());
    std::vector<gemmi::ReflnBlock> blocks = gemmi::as_refln_blocks(std::move(doc.blocks));
    gemmi::Logger logger;
    std::string why = name + " holds no merged reflection block";
    for (gemmi::ReflnBlock &rb : blocks) {
        if (rb.refln_loop == nullptr)
            continue;
        gemmi::CifToMtz converter;
        converter.spec_lines = SfCifSpec(*rb.refln_loop);
        try {
            const gemmi::Mtz mtz = converter.convert_block_to_mtz(rb, logger);
            RugnuxReflections out = RugnuxReflectionsFromMtz(mtz, name);
            out.source = "mmCIF block " + rb.block.name;
            return out;
        } catch (const std::exception &e) {
            why = e.what();
        }
    }
    throw std::runtime_error(why);
}

RugnuxReflections ReadRugnuxReflectionFile(const std::string &path) {
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw std::runtime_error("cannot open " + path);
    const std::string bytes{std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()};
    return ParseRugnuxReflections(bytes, path);
}