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

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <cmath>
#include <cstring>
#include <random>
#include <string>
#include <vector>

#include <gemmi/pdb.hpp>
#include <gemmi/solmask.hpp>
#include <gemmi/symmetry.hpp>

#include "../rugnux/ModelGrid.h"
#include "../rugnux/ModelMaskGPU.h"

namespace {

// The five groups of ModelValidation_ParallelGriddingMatchesGemmi: none, a centring, screws, a cubic
// body centring and a cubic face centring with a diamond glide.
const char *kCrysts[] = {
    "CRYST1   40.000   50.000   60.000  90.00  90.00  90.00 P 1           1\n",
    "CRYST1   40.000   50.000   60.000  90.00 100.00  90.00 C 1 2 1       4\n",
    "CRYST1   40.000   50.000   60.000  90.00  90.00  90.00 P 21 21 21    4\n",
    "CRYST1   60.000   60.000   60.000  90.00  90.00  90.00 I 2 3        24\n",
    "CRYST1   80.000   80.000   80.000  90.00  90.00  90.00 F 41 3 2     96\n",
};

// ModelValidationTest.cpp's ClusterPdb(): random atoms in one corner of the cell, from a fixed seed, with
// every fifth atom a nitrogen and every seventh an oxygen so that more than one radius is in play. A
// cluster this loose leaves many small pockets of solvent between its atoms: the islands.
gemmi::Structure ClusterModel(const char *cryst) {
    std::string pdb = cryst;
    std::mt19937 rng(20260902);
    std::uniform_real_distribution<double> x(2, 14), y(2, 16), z(2, 18);
    char line[96];
    for (int i = 1; i <= 150; i++) {
        const char *el = i % 5 == 0 ? "N" : i % 7 == 0 ? "O" : "C";
        std::snprintf(line, sizeof line,
                      "ATOM  %5d  %s   UNK A   1    %8.3f%8.3f%8.3f  1.00 20.00           %s\n",
                      i, el, x(rng), y(rng), z(rng), el);
        pdb += line;
    }
    pdb += "END\n";
    gemmi::Structure st = gemmi::read_pdb_string(pdb, "cluster");
    return st;
}

gemmi::Grid<float> MaskGrid(const gemmi::Structure &st, double d_min) {
    gemmi::Grid<float> grid;
    grid.unit_cell = st.cell;
    grid.spacegroup = st.find_spacegroup();
    grid.set_size_from_spacing(d_min / (2 * 1.5), gemmi::GridSizeRounding::Up);
    return grid;
}

// What the rigid body hands ModelMaskGPU: the atoms gemmi masks, fractional and wrapped, with gemmi's
// Refmac radius; every symmetry operator with every centring vector.
std::vector<ModelMaskAtom> MaskAtoms(const gemmi::Structure &st) {
    const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
    std::vector<ModelMaskAtom> atoms;
    for (const gemmi::Chain &ch : st.models[0].chains)
        for (const gemmi::Residue &res : ch.residues)
            for (const gemmi::Atom &a : res.atoms) {
                if (a.is_hydrogen() || a.occ <= 0)
                    continue;
                const gemmi::Fractional f = st.cell.fractionalize(a.pos).wrap_to_unit();
                const double r = masker.constant_r + masker.rprobe
                                 + gemmi::refmac_radius_for_bulk_solvent(a.element.elem);
                atoms.push_back(ModelMaskAtom{f.x, f.y, f.z, r});
            }
    return atoms;
}

std::vector<ModelMaskOp> MaskOps(const gemmi::SpaceGroup &sg) {
    const gemmi::GroupOps gops = sg.operations();
    std::vector<ModelMaskOp> ops;
    for (const auto &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;
            }
            ops.push_back(m);
        }
    return ops;
}

ModelMaskGrid MaskGridDescription(const gemmi::Grid<float> &grid) {
    ModelMaskGrid g;
    g.nu = grid.nu;
    g.nv = grid.nv;
    g.nw = grid.nw;
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++)
            g.orth[3 * i + j] = grid.unit_cell.orth.mat[i][j];
    g.volume = grid.unit_cell.volume;
    return g;
}

// |distance - radius| to the nearest image of any atom that reaches point (u, v, w), in A: how close
// the point lies to the edge of the mask.
double DistanceToMaskEdge(const gemmi::Grid<float> &grid, const std::vector<ModelMaskAtom> &atoms,
                          const std::vector<ModelMaskOp> &ops, int u, int v, int w) {
    const double p[3] = {static_cast<double>(u) / grid.nu, static_cast<double>(v) / grid.nv,
                         static_cast<double>(w) / grid.nw};
    double best = INFINITY;
    for (const ModelMaskAtom &a : atoms)
        for (const ModelMaskOp &op : ops) {
            const double f[3] = {a.x, a.y, a.z};
            gemmi::Fractional d;
            for (int k = 0; k < 3; k++) {
                const double x = op.rot[3 * k] * f[0] + op.rot[3 * k + 1] * f[1] + op.rot[3 * k + 2] * f[2]
                                 + op.tran[k];
                d.at(k) = p[k] - x - std::round(p[k] - x);
            }
            best = std::min(best, std::fabs(grid.unit_cell.orthogonalize_difference(d).length() - a.radius));
        }
    return best;
}

}  // namespace

// The whole mask against gemmi's put_mask_on_grid() - and so against the CPU path's PutMaskOnGrid() -
// the island step alone against gemmi's on gemmi's own pre-island mask (exact), and three repeats bit
// for bit. The masking itself may differ from gemmi only at points that sit at an atom's radius, where
// the float distance rounds the other way.
TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") {
    if (get_gpu_count() == 0)
        SKIP("No GPU");
    CudaStream stream;
    bool any_islands = false, any_odd_size = false;
    for (const char *cryst : kCrysts) {
        const gemmi::Structure st = ClusterModel(cryst);
        const gemmi::SpaceGroup *sg = st.find_spacegroup();
        REQUIRE(sg != nullptr);
        const std::vector<ModelMaskAtom> atoms = MaskAtoms(st);
        const std::vector<ModelMaskOp> ops = MaskOps(*sg);
        CudaDevicePtr<ModelMaskAtom> atoms_d(atoms.size());
        REQUIRE(cudaMemcpyAsync(atoms_d, atoms.data(), atoms.size() * sizeof(ModelMaskAtom), cudaMemcpyHostToDevice, stream) == cudaSuccess);

        for (double d_min : {6.0, 3.5}) {
            const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
            gemmi::Grid<float> pre = MaskGrid(st, d_min);
            masker.clear(pre);
            masker.mask_points(pre, st.models[0]);
            masker.symmetrize(pre);
            gemmi::Grid<float> ref = pre;
            const int islands = masker.remove_islands(ref);
            masker.shrink(ref);
            const size_t n = ref.data.size();
            // The rigid body's CPU path masks with PutMaskOnGrid, which is gemmi's put_mask_on_grid bit
            // for bit - checked here on this grid, so what the GPU is held to below is the CPU path too.
            gemmi::Grid<float> cpu = MaskGrid(st, d_min);
            PutMaskOnGrid(cpu, st.models[0], OrbitLeaders(cpu, 4), 4);
            REQUIRE(std::memcmp(cpu.data.data(), ref.data.data(), n * sizeof(float)) == 0);
            any_islands |= islands > 0;
            any_odd_size |= ref.nu % 8 != 0 || ref.nv % 8 != 0 || ref.nw % 8 != 0;

            ModelMaskGPU mask(stream, n);
            mask.SetGrid(MaskGridDescription(ref), ops);
            CudaDevicePtr<float> mask_d(n);
            std::vector<float> out(n), first;

            // The island step alone, on gemmi's pre-island mask: exact.
            REQUIRE(cudaMemcpyAsync(mask_d, pre.data.data(), n * sizeof(float), cudaMemcpyHostToDevice, stream) == cudaSuccess);
            mask.RemoveIslands(mask_d);
            REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
            REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
            CHECK(std::memcmp(out.data(), ref.data.data(), n * sizeof(float)) == 0);

            for (int repeat = 0; repeat < 3; repeat++) {
                mask.Compute(atoms_d, static_cast<int>(atoms.size()), mask_d);
                REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
                REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
                if (repeat == 0) {
                    first = out;
                    size_t differ = 0;
                    for (size_t i = 0; i < n; i++)
                        if (out[i] != ref.data[i]) {
                            differ++;
                            const int u = static_cast<int>(i % ref.nu), v = static_cast<int>(i / ref.nu % ref.nv),
                                    w = static_cast<int>(i / (static_cast<size_t>(ref.nu) * ref.nv));
                            CHECK(DistanceToMaskEdge(ref, atoms, ops, u, v, w) < 1e-4);
                        }
                    INFO(sg->hm << " d " << d_min << " grid " << ref.nu << "x" << ref.nv << "x" << ref.nw
                                << " islands " << islands << ": " << differ << " points differ");
                    CHECK(differ <= 15 * n / 10000000 + 1);
                } else
                    CHECK(std::memcmp(out.data(), first.data(), n * sizeof(float)) == 0);
            }
        }
    }
    CHECK(any_islands);
    CHECK(any_odd_size);
}

// The shrink is not implemented: a grid fine enough for it to change anything is refused.
TEST_CASE("ModelMaskGPU_RefusesAGridTheShrinkWouldChange", "[ModelValidation][gpu]") {
    if (get_gpu_count() == 0)
        SKIP("No GPU");
    CudaStream stream;
    const gemmi::Structure st = ClusterModel(kCrysts[0]);
    const std::vector<ModelMaskOp> ops = MaskOps(*st.find_spacegroup());
    const gemmi::Grid<float> coarse = MaskGrid(st, 3.5);
    const gemmi::Grid<float> fine = MaskGrid(st, 2.0);
    ModelMaskGPU mask(stream, fine.point_count());
    CHECK_NOTHROW(mask.SetGrid(MaskGridDescription(coarse), ops));
    CHECK_THROWS(mask.SetGrid(MaskGridDescription(fine), ops));
}

#endif