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

#include "LoadFCalcFromMtz.h"

#include <cmath>
#include <cstdint>
#include <string>
#include <vector>

#include <gemmi/mtz.hpp>

#include "../common/Reflection.h"

std::vector<MergedReflection> LoadFCalcFromMtz(const std::string& path) {
    gemmi::Mtz mtz;
    mtz.read_file_gz(path, true);

    // Prefer a calculated structure factor F-model (e.g. a reference structure): I_ref = F^2.
    const gemmi::Mtz::Column* col = mtz.column_with_label("F-model", nullptr, 'F');
    const bool square = (col != nullptr);

    // Otherwise fall back to a merged/observed intensity column, used directly as I_ref. This
    // lets PixelRefine be SELF-SEEDED from the data's own previous merge (a jfjoch merge writes
    // IMEAN) instead of an external reference structure - the EM-style outer loop, free of the
    // reference's non-isomorphism bias.
    if (col == nullptr) {
        for (const char* label : {"IMEAN", "I", "IOBS", "Iobs", "I-obs"}) {
            col = mtz.column_with_label(label, nullptr, 'J');
            if (col != nullptr)
                break;
        }
    }
    if (col == nullptr) {
        for (const auto& c : mtz.columns)
            if (c.type == 'J') { col = &c; break; } // any mean-intensity column
    }
    if (col == nullptr)
        throw std::runtime_error("MTZ has no F-model or intensity (J) column to use as reference");

    std::vector<MergedReflection> result;
    result.reserve(static_cast<std::size_t>(mtz.nreflections));

    const std::size_t stride = mtz.columns.size();

    for (int i = 0; i < mtz.nreflections; ++i) {
        const std::size_t row = static_cast<std::size_t>(i) * stride;
        const float v = (*col)[static_cast<std::size_t>(i)];

        if (std::isnan(v))
            continue;

        MergedReflection r;
        r.h = static_cast<int32_t>(mtz.data[row + 0]);
        r.k = static_cast<int32_t>(mtz.data[row + 1]);
        r.l = static_cast<int32_t>(mtz.data[row + 2]);
        r.I = square ? v * v : v;
        r.sigma = NAN;
        r.d = 0.0f;

        result.emplace_back(r);
    }

    return result;
}