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

#include "RefinementGPU.h"

#include <algorithm>
#include <cmath>
#include <optional>

#include "gemmi/dencalc.hpp"
#include "gemmi/it92.hpp"
#include "gemmi/solmask.hpp"   // SolventMasker, refmac_radius_for_bulk_solvent

#include "../../common/ParallelFor.h"
#include "RefinementGPUEngine.h"

namespace {

using Table = gemmi::IT92<float>;

// One atom's density as the gather takes it, isotropic, with `blur`.
ModelDensityAtom DensityAtom(const gemmi::Atom &atom, double b_iso, const gemmi::DensityCalculator<Table, float> &dc) {
    using CReal = Table::Coef::coef_type;
    const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
    const CReal b = static_cast<CReal>(b_iso + dc.blur);
    const auto precal = coef.precalculate_density_iso(b, 0.0f);
    const double radius = dc.estimate_radius(precal, b);
    ModelDensityAtom a{};
    a.occ = atom.occ;
    a.aniso = 0;
    a.radius = static_cast<float>(radius);
    a.max_r2 = static_cast<float>(radius * radius);
    for (int k = 0; k < 5; k++) {
        a.a[k] = precal.a[k];
        a.b[k][0] = precal.b[k];
    }
    return a;
}

std::vector<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
    std::vector<const gemmi::Atom *> atoms;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                atoms.push_back(&a);
    return atoms;
}

class RefinementEngineGPU : public RefinementEngine {
public:
    RefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                        const RefinementReflections &refl, const gemmi::Model &model, double b_max, size_t nthreads)
        : device_(device), cell_(cell), sg_(sg), d_min_(d_min), nthreads_(nthreads) {
        const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, 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;
        setup_.centring = static_cast<double>(gops.cen_ops.size());
        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);
            }
        for (size_t m = 0; m < refl.Size(); m++) {
            setup_.rows.push_back({refl.hkl[m][0], refl.hkl[m][1], refl.hkl[m][2]});
            setup_.inv_d2.push_back(refl.inv_d2[m]);
            setup_.fo.push_back(refl.fo[m]);
            setup_.work.push_back(!refl.free[m]);
            setup_.centric.push_back(refl.centric[m]);
        }

        // The adjoint of the composition on cuFFT's half-u spectrum: the term of row m and operator o, at
        // k = hR, goes to k where k_u > 0, conjugated to -k where k_u < 0, and to both on the planes k_u = 0
        // and k_u = nu/2, which hold their own Friedel mates.
        const int nu = grid.nu, nv = grid.nv, nw = grid.nw;
        auto point = [&](int ku, int kv, int kw) {
            return ku + (nu / 2 + 1) * (gemmi::modulo(kv, nv) + nv * gemmi::modulo(kw, nw));
        };
        std::vector<std::pair<int, int>> entries;
        const int n_ops = static_cast<int>(gops.sym_ops.size());
        for (size_t m = 0; m < refl.Size(); m++)
            for (int o = 0; o < n_ops; o++) {
                const gemmi::Miller k = gops.sym_ops[o].apply_to_hkl(refl.hkl[m]);
                const int term = (static_cast<int>(m) * n_ops + o) * 2;
                if (k[0] > 0 && 2 * k[0] != nu)
                    entries.emplace_back(point(k[0], k[1], k[2]), term);
                else if (k[0] < 0 && -2 * k[0] != nu)
                    entries.emplace_back(point(-k[0], -k[1], -k[2]), term + 1);
                else {
                    const int ku = std::abs(k[0]);
                    entries.emplace_back(point(ku, k[1], k[2]), term);
                    entries.emplace_back(point(ku, -k[1], -k[2]), term + 1);
                }
            }
        std::stable_sort(entries.begin(), entries.end(),
                         [](const auto &a, const auto &b) { return a.first < b.first; });
        for (size_t e = 0; e < entries.size(); e++) {
            if (e == 0 || entries[e].first != entries[e - 1].first) {
                setup_.scatter_point.push_back(entries[e].first);
                setup_.scatter_start.push_back(static_cast<int>(e));
            }
            setup_.scatter_term.push_back(entries[e].second);
        }
        setup_.scatter_start.push_back(static_cast<int>(entries.size()));

        // The gather's reserve for every atom at b_max with the largest blur a refinement sets (B down to 0).
        dc_ = Calculator(RefinementBlur(cell, sg, d_min, 0.0));
        std::vector<ModelDensityAtom> widest;
        for (const gemmi::Atom *a : Atoms(model))
            widest.push_back(DensityAtom(*a, b_max, dc_));
        setup_.max_atoms = widest.size();
        setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, widest);
        supported_ = !setup_.rows.empty() && ModelDensityGPU::Supports(setup_.grid, widest) &&
                     setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS;
        if (supported_) {
            fft_work_ = RefinementGPUEngine::FFTWorkBytes(device, setup_.grid);
            // the restraints: about thirty records and slots per atom
            bytes_ = RefinementGPUEngine::DeviceBytes(setup_, fft_work_, 30 * widest.size(), 60 * widest.size());
        }
    }

    bool Supported() const { return supported_; }
    size_t Bytes() const { return bytes_; }
    void Reserve() { engine_ = std::make_unique<RefinementGPUEngine>(device_, setup_, fft_work_); }

    void SetBlur(double blur) override {
        dc_ = Calculator(blur);
        last_b_.clear();
        engine_->SetBlur(blur);
    }

    void SetRowModel(const XrayRowModel &rows) override {
        std::vector<std::array<float, 2>> fmask(rows.fmask.size());
        for (size_t m = 0; m < fmask.size(); m++)
            fmask[m] = {rows.fmask[m].real(), rows.fmask[m].imag()};
        engine_->SetRowModel(rows.k_total, fmask, rows.d, rows.var);
    }

    bool OnGPU() const override { return true; }

    void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
        SetAtoms(model);
        std::vector<std::array<float, 2>> out;
        engine_->Fcalc(out);
        f.resize(out.size());
        for (size_t m = 0; m < out.size(); m++)
            f[m] = {out[m][0], out[m][1]};
    }

    void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
        const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
        std::vector<ModelMaskAtom> mask_atoms;
        for (const gemmi::Atom *a : Atoms(model)) {
            if ((masker.ignore_hydrogen && a->is_hydrogen()) || (masker.ignore_zero_occupancy_atoms && a->occ <= 0))
                continue;
            const gemmi::Fractional f0 = cell_.fractionalize(a->pos);
            double frac[3] = {f0.x, f0.y, f0.z};
            for (double &x : frac) {
                x -= std::floor(x);
                if (x >= 1.0)
                    x = 0.0;
            }
            mask_atoms.push_back({frac[0], frac[1], frac[2],
                                  masker.constant_r + masker.rprobe +
                                      gemmi::refmac_radius_for_bulk_solvent(a->element.elem)});
        }
        std::vector<std::array<float, 2>> out;
        engine_->Fmask(mask_atoms, out);
        f.resize(out.size());
        for (size_t m = 0; m < out.size(); m++)
            f[m] = {out[m][0], out[m][1]};
    }

    double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
        SetAtoms(model);
        std::vector<std::array<float, 4>> g;
        double value = 0;
        {
            const RefinementStage stage("gpu.evaluate");
            value = engine_->Evaluate(grad != nullptr ? &g : nullptr);
        }
        if (grad != nullptr) {
            grad->resize(g.size());
            for (size_t i = 0; i < g.size(); i++)
                (*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
        }
        return value;
    }

    double EvaluateWithRestraints(const gemmi::Model &model, int which, std::vector<std::array<double, 4>> *grad,
                                  std::vector<std::array<double, 4>> *restraint_grad, double &restraints) override {
        SetAtoms(model);
        std::vector<std::array<float, 4>> g, gr;
        double value = 0;
        {
            const RefinementStage stage("gpu.evaluate");
            value = engine_->EvaluateWithRestraints(grad != nullptr ? &g : nullptr, which, [&](std::array<float, 4> *x) {
                for (size_t i = 0; i < atom_ptrs_.size(); i++) {
                    const gemmi::Atom &a = *atom_ptrs_[i];
                    x[i] = {static_cast<float>(a.pos.x), static_cast<float>(a.pos.y), static_cast<float>(a.pos.z), a.b_iso};
                }
            }, restraint_grad != nullptr ? &gr : nullptr, restraints);
        }
        if (grad != nullptr) {
            grad->resize(g.size());
            for (size_t i = 0; i < g.size(); i++)
                (*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
        }
        if (restraint_grad != nullptr) {
            restraint_grad->resize(gr.size());
            for (size_t i = 0; i < gr.size(); i++)
                (*restraint_grad)[i] = {gr[i][0], gr[i][1], gr[i][2], gr[i][3]};
        }
        return value;
    }

    void SetRestraints(int which, const RestraintSet &set) override {
        engine_->SetRestraints(which, set.records, set.plane_atoms, set.slot_start, set.slot_index, set.n_slots);
    }

    double EvaluateRestraints(int which, const std::vector<std::array<double, 4>> &xyzb,
                              std::vector<std::array<double, 4>> *grad) override {
        std::vector<std::array<float, 4>> x(xyzb.size()), g;
        for (size_t i = 0; i < xyzb.size(); i++)
            x[i] = {static_cast<float>(xyzb[i][0]), static_cast<float>(xyzb[i][1]), static_cast<float>(xyzb[i][2]),
                    static_cast<float>(xyzb[i][3])};
        double value = 0;
        {
            const RefinementStage stage("gpu.restraints");
            value = engine_->EvaluateRestraints(which, x, grad != nullptr ? &g : nullptr);
        }
        if (grad != nullptr) {
            grad->resize(g.size());
            for (size_t i = 0; i < g.size(); i++)
                (*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
        }
        return value;
    }

private:
    gemmi::DensityCalculator<Table, float> Calculator(double blur) const {
        gemmi::DensityCalculator<Table, float> dc;
        dc.d_min = d_min_;
        dc.rate = 1.5;
        dc.grid.unit_cell = cell_;
        dc.grid.spacegroup = &sg_;
        dc.blur = blur;
        return dc;
    }

    // The atoms' grid positions to the device, and their densities where some B moved (made again only for the
    // atoms whose B moved).
    void SetAtoms(const gemmi::Model &model) {
        std::optional<RefinementStage> stage("gpu.atoms_host");
        atom_ptrs_ = Atoms(model);
        const std::vector<const gemmi::Atom *> &atoms = atom_ptrs_;
        const int n = static_cast<int>(atoms.size());
        if (last_b_.size() != atoms.size()) {
            last_b_.assign(atoms.size(), NAN);
            atoms_.resize(atoms.size());
        }
        std::vector<std::array<float, 4>> pos(atoms.size());
        std::vector<char> changed(atoms.size(), 0);
        const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
        ParallelChunks(n, ThreadsForWork(atoms.size(), nthreads_, 4096), [&](int lo, int hi) {
            for (int i = lo; i < hi; i++) {
                if (!(atoms[i]->b_iso == last_b_[i])) {
                    atoms_[i] = DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
                    last_b_[i] = atoms[i]->b_iso;
                    changed[i] = 1;
                }
                const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos);
                pos[i] = {static_cast<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(f.z * n3[2]), 0};
            }
        });
        stage.reset();
        const RefinementStage upload("gpu.upload");
        if (std::find(changed.begin(), changed.end(), 1) != changed.end())
            engine_->SetDensities(atoms_);
        engine_->SetPositions(pos);
    }

    int device_;
    gemmi::UnitCell cell_;
    const gemmi::SpaceGroup &sg_;
    double d_min_;
    size_t nthreads_;
    RefinementGPUSetup setup_;
    gemmi::DensityCalculator<Table, float> dc_;
    bool supported_ = false;
    size_t fft_work_ = 0, bytes_ = 0;
    std::unique_ptr<RefinementGPUEngine> engine_;
    std::vector<const gemmi::Atom *> atom_ptrs_;   // the model's atoms at the last SetAtoms()
    std::vector<float> last_b_;
    std::vector<ModelDensityAtom> atoms_;
};

}  // namespace

std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell,
                                                          const gemmi::SpaceGroup &sg, double d_min,
                                                          const RefinementReflections &refl,
                                                          const gemmi::Model &model, double b_max, size_t budget,
                                                          size_t *device_bytes, size_t nthreads) {
    auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max, nthreads);
    if (device_bytes != nullptr)
        *device_bytes = engine->Bytes();
    if (!engine->Supported() || engine->Bytes() > budget)
        return nullptr;
    engine->Reserve();
    return engine;
}