// 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 <map>
#include <memory>
#include <mutex>
#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 "ModelGrid.h"   // IonFormFactor, AnisoRadiusB
#include "RefinementGPUEngine.h"

namespace {

using Table = gemmi::IT92<float>;

// One atom's density as the gather takes it, anisotropic with U (as PutModelDensityOnGrid() grids it), with `blur`.
ModelDensityAtom AnisoDensityAtom(const gemmi::Atom &atom, const gemmi::SMat33<float> &u,
                                  const gemmi::DensityCalculator<Table, float> &dc) {
    using CReal = Table::Coef::coef_type;
    const auto &coef = IonFormFactor<float>(atom.element.elem, atom.charge);
    const auto aniso_b = u.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
    const CReal b_max = AnisoRadiusB(aniso_b);
    const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, 0.0f), b_max);
    const auto precal = coef.precalculate_density_aniso_b(aniso_b, 0.0f);
    ModelDensityAtom a{};
    a.occ = atom.occ;
    a.aniso = 1;
    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];
        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]);
    }
    return a;
}

// 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 = IonFormFactor<float>(atom.element.elem, atom.charge);
    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;
}

// Whether `atoms` is Atoms(model) still: as many atoms, the first and the last the same ones.
bool SameAtoms(const gemmi::Model &model, const std::vector<const gemmi::Atom *> &atoms) {
    size_t n = 0;
    const gemmi::Atom *first = nullptr, *last = nullptr;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            if (!r.atoms.empty()) {
                if (first == nullptr)
                    first = &r.atoms.front();
                last = &r.atoms.back();
                n += r.atoms.size();
            }
    return n > 0 && n == atoms.size() && first == atoms.front() && last == atoms.back();
}

// The engine's single-precision gradient into the caller's doubles, element by element in parallel: bit for bit
// the serial loop's.
void Widen(const std::array<float, 4> *from, size_t n, std::vector<std::array<double, 4>> &to, size_t nthreads) {
    to.resize(n);
    ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
        for (int i = lo; i < hi; i++)
            to[i] = {from[i][0], from[i][1], from[i][2], from[i][3]};
    });
}

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,
                        size_t restraint_records, size_t restraint_slots)
        : device_(device), cell_(cell), sg_(sg), d_min_(d_min), nthreads_(nthreads),
          setup_(std::make_shared<RefinementGPUSetup>()), cards_(std::make_shared<Cards>()) {
        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;
        // the mask on a grid of its own, as BulkSolventMask has it
        const gemmi::Grid<float> mask_grid = BulkSolventMask::MaskGrid(cell, sg, d_min);
        setup_->mask_grid = setup_->grid;
        setup_->mask_grid.nu = mask_grid.nu;
        setup_->mask_grid.nv = mask_grid.nv;
        setup_->mask_grid.nw = mask_grid.nw;
        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]);
        }
        // a twin law's mates not observed: structure factors only
        for (const gemmi::Miller &h : refl.mate_hkl) {
            setup_->rows.push_back({h[0], h[1], h[2]});
            setup_->inv_d2.push_back(cell.calculate_1_d2(h));
            setup_->fo.push_back(0);
            setup_->work.push_back(0);
            setup_->centric.push_back(0);
        }
        setup_->mate = refl.mate;

        // 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());
        const std::vector<gemmi::Miller> row_hkl = refl.RowHkl();
        for (size_t m = 0; m < row_hkl.size(); m++)
            for (int o = 0; o < n_ops; o++) {
                const gemmi::Miller k = gops.sym_ops[o].apply_to_hkl(row_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_);
            // the restraints, where they are not given: about thirty records and sixty slots per atom
            if (restraint_records == 0 && restraint_slots == 0) {
                restraint_records = 30 * widest.size();
                restraint_slots = 60 * widest.size();
            }
            bytes_.engine = RefinementGPUEngine::DeviceBytes(*setup_, restraint_records, restraint_slots);
            bytes_.card = RefinementGPUEngine::CardBytes(*setup_);
            bytes_.lane = RefinementGPUEngine::LaneBytes(*setup_, fft_work_);
        }
    }

    // One like `main` on `device`: the same setup, and the card's part shared with the engines of main's kind there.
    RefinementEngineGPU(const RefinementEngineGPU &main, int device, size_t nthreads, bool own_lane)
        : device_(device), cell_(main.cell_), sg_(main.sg_), d_min_(main.d_min_), nthreads_(nthreads), setup_(main.setup_),
          dc_(Calculator(RefinementBlur(cell_, sg_, d_min_, 0.0))), supported_(main.supported_),
          fft_work_(RefinementGPUEngine::FFTWorkBytes(device, *setup_)), bytes_(main.bytes_), own_lane_(own_lane), cards_(main.cards_) {}

    bool Supported() const { return supported_; }
    RefinementGPUBytes Bytes() const { return bytes_; }
    void Reserve() {
        engine_ = std::make_unique<RefinementGPUEngine>(device_, *setup_, cards_->Get(device_, *setup_), fft_work_, own_lane_);
    }

    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()};
        // per row of the engine: the mates' rows have no observation, and nothing of theirs is read but fmask
        std::vector<float> k_total = rows.k_total, d = rows.d, var = rows.var;
        k_total.resize(setup_->rows.size(), 0.0f);
        d.resize(setup_->rows.size(), 0.0f);
        var.resize(setup_->rows.size(), 1.0f);
        engine_->SetRowModel(k_total, fmask, d, var, rows.twin_fraction);
    }

    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);
        double value = 0;
        {
            const RefinementStage stage("gpu.evaluate");
            value = engine_->Evaluate(grad != nullptr);
        }
        if (grad != nullptr) {
            Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
            CopyUGradient();
        }
        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);
        double value = 0;
        {
            const RefinementStage stage("gpu.evaluate");
            value = engine_->EvaluateWithRestraints(grad != 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, restraints);
        }
        if (grad != nullptr) {
            Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
            CopyUGradient();
        }
        if (restraint_grad != nullptr)
            Widen(engine_->HostRestraintGradient(), atom_ptrs_.size(), *restraint_grad, nthreads_);
        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:
    // The U derivatives of the last evaluation, where some atom is anisotropic.
    void CopyUGradient() {
        if (!AnyAniso()) {
            grad_u_.clear();
            return;
        }
        const float *g = engine_->HostUGradient();
        grad_u_.resize(atom_ptrs_.size());
        for (size_t i = 0; i < atom_ptrs_.size(); i++)
            for (int c = 0; c < 6; c++)
                grad_u_[i][c] = g[6 * i + c];
    }

    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;
        dc.cutoff = REFINEMENT_DENSITY_CUTOFF;
        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");
        // the pointers kept from the last call where the model is the same one still (its context's own)
        if (!SameAtoms(model, atom_ptrs_))
            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);
            last_key_.assign(atoms.size(), AtomKey());
            last_aniso_.assign(atoms.size(), AtomAniso());
            atoms_.resize(atoms.size());
            pos_.resize(atoms.size());
        }
        changed_.assign(atoms.size(), 0);
        const int n3[3] = {setup_->grid.nu, setup_->grid.nv, setup_->grid.nw};
        const std::vector<AtomAniso> *aniso = aniso_ != nullptr && aniso_->size() == atoms.size() ? aniso_ : nullptr;
        ParallelChunks(n, ThreadsForWork(atoms.size(), nthreads_, 4096), [&](int lo, int hi) {
            for (int i = lo; i < hi; i++) {
                // made again where its B or its anisotropy moved, or where the solvent build put it in, took it out or
                // relabelled it
                const AtomKey key{atoms[i]->occ, static_cast<int>(atoms[i]->element.elem), atoms[i]->charge};
                const AtomAniso none;
                const AtomAniso &an = aniso != nullptr ? (*aniso)[i] : none;
                if (!(atoms[i]->b_iso == last_b_[i]) || !(key == last_key_[i]) || an.on != last_aniso_[i].on ||
                    (an.on && an.dev != last_aniso_[i].dev)) {
                    atoms_[i] = an.on ? AnisoDensityAtom(*atoms[i], AnisoU(an, atoms[i]->b_iso), dc_)
                                      : DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
                    last_b_[i] = atoms[i]->b_iso;
                    last_key_[i] = key;
                    last_aniso_[i] = an;
                    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_.data(), atoms_.size());
        engine_->SetPositions(pos_.data(), pos_.size());
    }

    int device_;
    gemmi::UnitCell cell_;
    const gemmi::SpaceGroup &sg_;
    double d_min_;
    size_t nthreads_;
    std::shared_ptr<RefinementGPUSetup> setup_;   // read only once made: shared with the engines like this one
    gemmi::DensityCalculator<Table, float> dc_;
    bool supported_ = false;
    size_t fft_work_ = 0;
    RefinementGPUBytes bytes_;
    bool own_lane_ = true;
    // The cards' shared parts of the engines like this one, each kept while an engine on its card is.
    struct Cards {
        std::mutex m;
        std::map<int, std::weak_ptr<RefinementGPUCard>> card;
        std::shared_ptr<RefinementGPUCard> Get(int device, const RefinementGPUSetup &setup) {
            const std::lock_guard lock(m);
            std::shared_ptr<RefinementGPUCard> c = card[device].lock();
            if (!c) {
                c = RefinementGPUEngine::MakeCard(device, setup);
                card[device] = c;
            }
            return c;
        }
    };
    std::shared_ptr<Cards> cards_;
    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_;
    struct AtomKey {
        float occ = -1;
        int element = -1;
        signed char charge = 0;
        bool operator==(const AtomKey &o) const { return occ == o.occ && element == o.element && charge == o.charge; }
    };
    std::vector<AtomKey> last_key_;
    std::vector<AtomAniso> last_aniso_;
    PinnedVector<ModelDensityAtom> atoms_;   // pinned: a B rung sends them every evaluation (37 MB on a large model)
    PinnedVector<std::array<float, 4>> pos_;
    std::vector<char> changed_;
};

}  // 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,
                                                          RefinementGPUBytes *bytes, size_t nthreads,
                                                          size_t restraint_records, size_t restraint_slots) {
    auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max, nthreads,
                                                        restraint_records, restraint_slots);
    if (bytes != nullptr)
        *bytes = engine->Bytes();
    if (!engine->Supported() || engine->Bytes().engine + engine->Bytes().card + engine->Bytes().lane > budget)
        return nullptr;
    engine->Reserve();
    return engine;
}

std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(const RefinementEngine &like, int device, size_t nthreads,
                                                          bool own_lane) {
    auto engine = std::make_unique<RefinementEngineGPU>(dynamic_cast<const RefinementEngineGPU &>(like), device, nthreads,
                                                        own_lane);
    engine->Reserve();
    return engine;
}