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

#include "RugnuxStructure.h"

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstring>

#include "gemmi/asudata.hpp"
#include "gemmi/calculate.hpp"   // calculate_fractional_box
#include "gemmi/ccp4.hpp"
#include "gemmi/fourier.hpp"     // get_size_for_hkl, get_f_phi_on_grid

#include "../image_analysis/structure_refinement/ModelFFT.h"
#include "../writer/WriteModel.h"

namespace {
    const MtzTable::Column &FindColumn(const MtzTable &mtz, const std::string &label) {
        for (const auto &c: mtz.columns)
            if (c.label == label)
                return c;
        throw std::runtime_error("map coefficients have no column " + label);
    }
}

RugnuxStructure::RugnuxStructure(std::shared_ptr<const gemmi::Structure> model, const UnitCell &cell,
                                 const gemmi::SpaceGroup &space_group, std::optional<MtzTable> coefficients)
    : model(std::move(model)), cell(cell), space_group(space_group), coefficients(std::move(coefficients)) {}

bool RugnuxStructure::HasMaps() const {
    return coefficients.has_value();
}

bool RugnuxStructure::HasAnomalousMap() const {
    if (!coefficients)
        return false;
    for (const auto &c: coefficients->columns)
        if (c.label == "ANOM")
            return true;
    return false;
}

std::string RugnuxStructure::ModelCif() const {
    std::unique_lock ul(m);
    if (!model_cif)
        model_cif = PlacedModelMmcif(*model, cell, space_group);
    return *model_cif;
}

std::string RugnuxStructure::Map(RugnuxMapType type) const {
    if (!coefficients || (type == RugnuxMapType::Anomalous && !HasAnomalousMap()))
        return {};
    std::unique_lock ul(m);
    std::optional<std::string> &cached = type == RugnuxMapType::TwoFoFc ? map_2fofc
                                         : type == RugnuxMapType::FoFc ? map_fofc : map_anom;
    if (cached)
        return *cached;

    const MtzTable &mtz = *coefficients;
    const auto &h = FindColumn(mtz, "H"), &k = FindColumn(mtz, "K"), &l = FindColumn(mtz, "L");
    const char *f_label = type == RugnuxMapType::TwoFoFc ? "FWT" : type == RugnuxMapType::FoFc ? "DELFWT" : "ANOM";
    const char *phi_label = type == RugnuxMapType::TwoFoFc ? "PHWT" : type == RugnuxMapType::FoFc ? "PHDELWT" : "PHANOM";
    const auto &f = FindColumn(mtz, f_label);
    const auto &phi = FindColumn(mtz, phi_label);

    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(mtz.space_group_name);
    if (sg == nullptr)
        throw std::runtime_error("unknown space group '" + mtz.space_group_name + "' of the map coefficients");
    gemmi::AsuData<std::complex<float>> coef;
    coef.unit_cell_ = mtz.unit_cell;
    coef.spacegroup_ = sg;
    for (size_t i = 0; i < h.values.size(); i++)
        if (std::isfinite(f.values[i]) && std::isfinite(phi.values[i]))   // missing: no Bijvoet pair (ANOM)
            coef.v.push_back({{static_cast<int>(h.values[i]), static_cast<int>(k.values[i]), static_cast<int>(l.values[i])},
                          std::polar(f.values[i], phi.values[i] * static_cast<float>(M_PI / 180.0))});
    coef.ensure_sorted();
    // Sampled as ModelValidation samples its maps.
    const std::array<int, 3> size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0);

    gemmi::Ccp4<float> map;
    map.grid = MapFromFPhi(gemmi::get_f_phi_on_grid<float>(coef, size, true));
    map.update_ccp4_header(2, true);   // the statistics of the whole cell

    gemmi::Structure st = *model;
    st.cell = mtz.unit_cell;
    const gemmi::Box<gemmi::Fractional> model_box = gemmi::calculate_fractional_box(st, MAP_MARGIN_A);
    gemmi::Box<gemmi::Fractional> box;
    for (int i = 0; i < 3; i++) {
        const double centre = (model_box.minimum.at(i) + model_box.maximum.at(i)) / 2;
        const double half = std::max(0.5, (model_box.maximum.at(i) - model_box.minimum.at(i)) / 2);
        box.minimum.at(i) = centre - half;
        box.maximum.at(i) = centre + half - 1e-9;   // one cell is exactly one cell of grid points
    }
    map.set_extent(box);

    std::string out(4 * map.ccp4_header.size() + sizeof(float) * map.grid.data.size(), '\0');
    memcpy(out.data(), map.ccp4_header.data(), 4 * map.ccp4_header.size());
    memcpy(out.data() + 4 * map.ccp4_header.size(), map.grid.data.data(), sizeof(float) * map.grid.data.size());
    cached = std::move(out);
    return *cached;
}