// 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 <array>
#include <cstring>
#include <random>
#include <vector>

#include <gemmi/dencalc.hpp>
#include <gemmi/it92.hpp>

#include "../image_analysis/structure_refinement/ModelDensityGPU.h"
#include "../image_analysis/structure_refinement/ModelGrid.h"
#include "../image_analysis/structure_refinement/ModelStructureFactorsGPU.h"

// The CUDA gather is the CPU's density walk, bit for bit (ModelDensityGPU.h): PutModelDensityOnGrid() with no
// orbits to symmetrize - one copy of the content, which is what the gather puts on the grid - against
// ModelDensityGPU fed as ModelStructureFactorsGPU feeds it. The grids are compared as bit patterns, as in
// ModelDensityMetalTest.cpp.

namespace {
    // Atoms scattered over the whole cell and a half cell beyond it on every side - so that boxes wrap at every
    // face and centres lie outside the cell - of a few elements, a quarter of them anisotropic, with occupancies
    // below one.
    gemmi::Structure random_crystal(const gemmi::UnitCell &cell, const char *group) {
        std::mt19937 rng(11);
        std::uniform_real_distribution<double> frac(-0.5, 1.5), b(8.0, 60.0), occ(0.3, 1.0), u(0.05, 0.4);
        const gemmi::El elements[] = {gemmi::El::C, gemmi::El::N, gemmi::El::O, gemmi::El::S};
        gemmi::Structure st;
        st.cell = cell;
        st.spacegroup_hm = group;
        st.setup_cell_images();
        gemmi::Model model(1);
        gemmi::Chain chain("A");
        gemmi::Residue res;
        res.name = "UNK";
        res.seqid = gemmi::SeqId(1, ' ');
        for (int i = 0; i < 3000; i++) {
            gemmi::Atom a;
            a.name = "X";
            a.element = elements[i % 4];
            a.pos = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
            a.b_iso = static_cast<float>(b(rng));
            a.occ = static_cast<float>(occ(rng));
            if (i % 4 == 3) {
                const float d1 = static_cast<float>(u(rng)), d2 = static_cast<float>(u(rng)),
                            d3 = static_cast<float>(u(rng));
                a.aniso = {d1, d2, d3, 0.3f * std::min(d1, d2), -0.2f * std::min(d1, d3), 0.1f * std::min(d2, d3)};
            }
            res.atoms.push_back(a);
        }
        chain.residues.push_back(res);
        model.chains.push_back(chain);
        st.models.push_back(model);
        return st;
    }

    void check_same_grid(const gemmi::UnitCell &cell, const char *group, double d_min) {
        const gemmi::Structure st = random_crystal(cell, group);
        gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
        dc.d_min = d_min;
        dc.rate = 1.5;
        dc.set_grid_cell_and_spacegroup(st);
        dc.set_refmac_compatible_blur(st.models[0]);
        PutModelDensityOnGrid(dc, st.models[0], {}, 4, false);
        const gemmi::Grid<float> &cpu = dc.grid;

        ModelDensityGrid grid;
        grid.nu = cpu.nu;
        grid.nv = cpu.nv;
        grid.nw = cpu.nw;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++) {
                grid.orth[3 * i + j] = cell.orth.mat[i][j];
                grid.frac[3 * i + j] = cell.frac.mat[i][j];
            }
        std::vector<ModelDensityAtom> atoms;
        std::vector<int> mask_atom;
        std::vector<float> mask_radius;
        ModelDensityAtoms(st.models[0], dc, atoms, mask_atom, mask_radius);
        REQUIRE(ModelDensityGPU::Supports(grid, atoms));
        // ModelStructureFactorsGPU::Compute()'s centres: the fractional position times the grid size.
        std::vector<std::array<float, 4>> pos;
        for (const gemmi::Atom &a : st.models[0].chains[0].residues[0].atoms) {
            const gemmi::Fractional f = cell.fractionalize(a.pos);
            pos.push_back({static_cast<float>(f.x * grid.nu), static_cast<float>(f.y * grid.nv),
                           static_cast<float>(f.z * grid.nw), 0.0f});
        }

        CudaStream stream;
        ModelDensityGPU density(stream, atoms.size(), ModelDensityGPU::PairBound(grid, atoms),
                                ModelDensityGPU::Bricks(grid.nu, grid.nv, grid.nw));
        density.SetAtoms(grid, atoms);
        const size_t n = cpu.data.size();
        CudaDevicePtr<float4> pos_d(pos.size());
        CudaDevicePtr<float> grid_d(n);
        REQUIRE(cudaMemcpyAsync(pos_d, pos.data(), pos.size() * sizeof(float4), cudaMemcpyHostToDevice, stream) == cudaSuccess);
        density.Compute(pos_d, grid_d);
        std::vector<float> gpu(n);
        REQUIRE(cudaMemcpyAsync(gpu.data(), grid_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
        REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
        size_t different = 0;
        for (size_t i = 0; i < n; i++)
            different += std::memcmp(&cpu.data[i], &gpu[i], sizeof(float)) != 0;
        CHECK(different == 0);
    }
}

TEST_CASE("ModelDensityGPU_BitIdenticalToCPU", "[ModelValidation][gpu]") {
    if (get_gpu_count() == 0)
        SKIP("No GPU");
    // Triclinic, so that every term of the orthogonalization is there; then an orthorhombic group's grid.
    check_same_grid(gemmi::UnitCell(61, 68, 77, 81, 96, 107), "P 1", 2.0);
    check_same_grid(gemmi::UnitCell(70, 80, 90, 90, 90, 90), "P 21 21 21", 2.0);
}

#endif