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

// The reflection list follows ReciprocalGrid::prepare_asu_data() of GEMMI's recgrid.hpp
// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/recgrid.hpp)
// (c) Global Phasing Ltd., Mozilla Public License Version 2.0

#include "ModelStructureFactorsGPU.h"

#include <algorithm>
#include <cmath>

#include "gemmi/fourier.hpp"   // get_size_for_hkl, get_f_phi_on_grid
#include "gemmi/solmask.hpp"   // SolventMasker, refmac_radius_for_bulk_solvent

#include "RigidBodyRefine.h"   // RigidBodyZoneGrid

namespace {

using Table = gemmi::IT92<float>;

// The density calculator model validation grids F_calc with: d_min, rate 1.5 and the blur gemmi's
// set_refmac_compatible_blur() gives this model.
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                  const gemmi::SpaceGroup &sg, double d_min) {
    gemmi::DensityCalculator<Table, float> dc;
    dc.d_min = d_min;
    dc.rate = 1.5;
    dc.grid.unit_cell = cell;
    dc.grid.spacegroup = &sg;
    dc.set_refmac_compatible_blur(model);
    return dc;
}

// The reflections prepare_asu_data(d_min) lists for the half-l transform of an (nu, nv, nw) grid, in its
// order (h, then k, then l ascending): in the group's reciprocal ASU, strictly inside d_min, inside the
// grid's Nyquist box, not systematically absent and not 000.
std::vector<gemmi::Miller> AsuReflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                                          int nu, int nv, int nw) {
    const gemmi::Miller lim = cell.get_hkl_limits(d_min);
    const int max_h = std::min((nu - 1) / 2, lim[0]);
    const int max_k = std::min((nv - 1) / 2, lim[1]);
    const int max_l = std::min(nw / 2, lim[2]);
    const double max_1_d2 = 1.0 / (d_min * d_min);
    const gemmi::ReciprocalAsu asu(&sg);
    const gemmi::GroupOps gops = sg.operations();
    std::vector<gemmi::Miller> rows;
    gemmi::Miller hkl;
    for (hkl[0] = -max_h; hkl[0] <= max_h; ++hkl[0])
        for (hkl[1] = -max_k; hkl[1] <= max_k; ++hkl[1])
            for (hkl[2] = -max_l; hkl[2] <= max_l; ++hkl[2])
                if (asu.is_in(hkl) && cell.calculate_1_d2(hkl) < max_1_d2 && !gops.is_systematically_absent(hkl) &&
                    !(hkl[0] == 0 && hkl[1] == 0 && hkl[2] == 0))
                    rows.push_back(hkl);
    return rows;
}

}  // namespace

void ModelDensityAtoms(const gemmi::Model &model, const gemmi::DensityCalculator<Table, float> &dc,
                       std::vector<ModelDensityAtom> &atoms, std::vector<int> &mask_atom,
                       std::vector<float> &mask_radius) {
    atoms.clear();
    mask_atom.clear();
    mask_radius.clear();
    const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
    int index = 0;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &atom : r.atoms) {
                using CReal = Table::Coef::coef_type;
                const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
                const float addend = dc.addends.get(atom.element);
                ModelDensityAtom a{};
                a.occ = atom.occ;
                a.aniso = atom.aniso.nonzero();
                if (!a.aniso) {
                    const CReal b = static_cast<CReal>(atom.b_iso + dc.blur);
                    const auto precal = coef.precalculate_density_iso(b, addend);
                    a.radius = dc.estimate_radius(precal, b);
                    for (int k = 0; k < 5; k++) {
                        a.a[k] = precal.a[k];
                        a.b[k][0] = precal.b[k];
                    }
                } else {
                    const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
                    const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33);
                    a.radius = static_cast<float>(dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max));
                    const auto precal = coef.precalculate_density_aniso_b(aniso_b, addend);
                    for (int k = 0; k < 5; k++) {
                        a.a[k] = precal.a[k];
                        const gemmi::SMat33<float> &m = precal.b[k];
                        const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
                        std::copy(e, e + 6, a.b[k]);
                    }
                }
                atoms.push_back(a);
                if (!((masker.ignore_hydrogen && atom.is_hydrogen()) ||
                      (masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) {
                    mask_atom.push_back(index);
                    mask_radius.push_back(static_cast<float>(
                        masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem)));
                }
                ++index;
            }
}

ModelStructureFactorsGPU::ModelStructureFactorsGPU(int device, const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                   const gemmi::SpaceGroup &sg, double d_min)
    : device_(device), cell_(cell), sg_(&sg), d_min_(d_min), dc_(Calculator(model, cell, sg, d_min)) {
    const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);   // DensityCalculator's for d_min
    setup_.grid.nu = grid.nu;
    setup_.grid.nv = grid.nv;
    setup_.grid.nw = grid.nw;
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++) {
            setup_.grid.orth[3 * i + j] = cell.orth.mat[i][j];
            setup_.grid.frac[3 * i + j] = cell.frac.mat[i][j];
        }
    setup_.volume = cell.volume;

    const gemmi::GroupOps gops = sg.operations();
    setup_.den = gemmi::Op::DEN;
    for (const gemmi::Op &op : gops.sym_ops) {
        std::array<int, 12> o{};
        for (int i = 0; i < 3; i++) {
            for (int j = 0; j < 3; j++)
                o[3 * i + j] = op.rot[i][j];
            o[9 + i] = op.tran[i];
        }
        setup_.sym_ops.push_back(o);
    }
    for (const gemmi::Op::Tran &cen : gops.cen_ops)
        for (const gemmi::Op &op : gops.sym_ops) {
            ModelMaskOp m{};
            for (int i = 0; i < 3; i++) {
                for (int j = 0; j < 3; j++)
                    m.rot[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
                m.tran[i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
            }
            setup_.mask_ops.push_back(m);
        }

    rows_ = AsuReflections(cell, sg, d_min, grid.nu, grid.nv, grid.nw);
    const double n_cen = static_cast<double>(gops.cen_ops.size());
    for (const gemmi::Miller &h : rows_) {
        setup_.rows.push_back({h[0], h[1], h[2]});
        // prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4)
        setup_.row_scale.push_back(n_cen * std::exp(dc_.blur * 0.25 * cell.calculate_1_d2(h)));
    }

    std::vector<ModelDensityAtom> atoms;
    std::vector<int> mask_atom;
    std::vector<float> mask_radius;
    ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius);
    setup_.max_atoms = atoms.size();
    setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, atoms);
    supported_ = !rows_.empty() && ModelDensityGPU::Supports(setup_.grid, atoms) &&
                 setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS;
    if (!supported_)
        return;

    // The largest map: Map()'s coefficients are on a subset of these reflections, and get_size_for_hkl()
    // only grows with the indices and the resolution it is given.
    gemmi::AsuData<std::complex<float>> all;
    all.unit_cell_ = cell;
    all.spacegroup_ = &sg;
    for (const gemmi::Miller &h : rows_)
        all.v.push_back({h, {1.0f, 0.0f}});
    map_size_ = gemmi::get_size_for_hkl(all, {{0, 0, 0}}, 3.0);
    // MapFromFPhi() transforms a half-l grid of nw / 2 + 1 planes back to 2 (nw / 2) of them.
    const int map_nw = 2 * (map_size_[2] / 2);
    setup_.map_points = static_cast<size_t>(map_size_[0]) * map_size_[1] * map_nw;
    setup_.map_complex_points = static_cast<size_t>(map_size_[0]) * map_size_[1] * (map_nw / 2 + 1);
    fft_work_bytes_ = ModelStructureFactorsGPUEngine::FFTWorkBytes(device_, setup_.grid, map_size_[0], map_size_[1], map_nw);
    device_bytes_ = ModelStructureFactorsGPUEngine::DeviceBytes(setup_, fft_work_bytes_);
}

ModelStructureFactorsGPU::~ModelStructureFactorsGPU() = default;

void ModelStructureFactorsGPU::Reserve() {
    std::lock_guard lock(m_);
    if (!engine_)
        engine_ = std::make_unique<ModelStructureFactorsGPUEngine>(device_, setup_, fft_work_bytes_);
}

void ModelStructureFactorsGPU::Compute(const gemmi::Model &model, gemmi::AsuData<std::complex<float>> &fcalc,
                                       gemmi::AsuData<std::complex<float>> &fmask) {
    std::vector<ModelDensityAtom> atoms;
    std::vector<int> mask_atom;
    std::vector<float> mask_radius;
    ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius);

    // Each atom in grid units and, for the mask, fractional - both wrapped into the cell.
    const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
    std::vector<std::array<float, 4>> pos;
    std::vector<std::array<double, 3>> frac;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms) {
                const gemmi::Fractional f0 = cell_.fractionalize(a.pos);
                std::array<double, 3> f{f0.x, f0.y, f0.z};
                std::array<float, 4> g{};
                for (int k = 0; k < 3; k++) {
                    f[k] -= std::floor(f[k]);
                    if (f[k] >= 1.0)   // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1
                        f[k] = 0.0;
                    g[k] = static_cast<float>(f[k] * n3[k]);
                    if (g[k] >= n3[k])
                        g[k] -= n3[k];
                }
                frac.push_back(f);
                pos.push_back(g);
            }
    std::vector<ModelMaskAtom> mask_atoms;
    for (size_t j = 0; j < mask_atom.size(); j++) {
        const std::array<double, 3> &f = frac[mask_atom[j]];
        mask_atoms.push_back({f[0], f[1], f[2], mask_radius[j]});
    }

    std::vector<std::array<float, 2>> fc, fm;
    {
        std::lock_guard lock(m_);
        engine_->Compute(atoms, pos, mask_atoms, fc, fm);
    }
    fcalc.v.clear();
    fmask.v.clear();
    fcalc.v.reserve(rows_.size());
    fmask.v.reserve(rows_.size());
    for (size_t i = 0; i < rows_.size(); i++) {
        fcalc.v.push_back({rows_[i], {fc[i][0], fc[i][1]}});
        fmask.v.push_back({rows_[i], {fm[i][0], fm[i][1]}});
    }
    fcalc.unit_cell_ = fmask.unit_cell_ = cell_;
    fcalc.spacegroup_ = fmask.spacegroup_ = sg_;
}

gemmi::Grid<float> ModelStructureFactorsGPU::Map(gemmi::AsuData<std::complex<float>> &coef) {
    coef.ensure_sorted();
    const std::array<int, 3> size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0);
    // ZYX: l fastest and halved, the layout a c2r transform takes, where gemmi's XYZ grid halves the slowest
    // axis. The coefficients written are the same.
    const gemmi::FPhiGrid<float> hkl = gemmi::get_f_phi_on_grid<float>(coef, size, true, gemmi::AxisOrder::ZYX);
    gemmi::Grid<float> map;
    map.spacegroup = coef.spacegroup_;
    map.unit_cell = coef.unit_cell_;
    // As MapFromFPhi(): 2 (nw - 1) planes back from nw stored ones.
    map.set_size(size[0], size[1], 2 * (hkl.nu - 1));
    map.axis_order = gemmi::AxisOrder::XYZ;
    std::lock_guard lock(m_);
    engine_->Map(map.nu, map.nv, map.nw, map.unit_cell.volume, reinterpret_cast<const float *>(hkl.data.data()),
                 map.data.data());
    return map;
}