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

#include <catch2/catch_all.hpp>

#ifdef __APPLE__

#include <cstring>
#include <random>
#include <vector>

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

#include "../image_analysis/structure_refinement/ModelDensityMetal.h"
#include "../image_analysis/structure_refinement/ModelGrid.h"

// The model density on the Apple GPU is the CPU's, bit for bit (ModelDensityMetal.h), so the two grids are compared
// as bit patterns. A Mac without a Metal device fails rather than skips - the macOS CI job is the only place this
// can be checked, as in BraggIntegrationEngineMetalTest.cpp.

namespace {
    // Atoms scattered over the whole cell - so that boxes wrap at every face - 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.0, 1.0), 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> cpu, gpu;
        for (auto *dc : {&cpu, &gpu}) {
            dc->d_min = d_min;
            dc->rate = 1.5;
            dc->set_grid_cell_and_spacegroup(st);
            dc->set_refmac_compatible_blur(st.models[0]);
        }
        // The orbits of the grid initialize_grid() sizes, as ModelValidation.cpp finds them.
        gemmi::Grid<float> sized;
        sized.unit_cell = cpu.grid.unit_cell;
        sized.spacegroup = cpu.grid.spacegroup;
        sized.set_size_from_spacing(cpu.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
        const std::vector<size_t> leaders = OrbitLeaders(sized, 4);
        PutModelDensityOnGrid(cpu, st.models[0], leaders, 4, false);
        const size_t before = ModelDensityMetalGrids();
        PutModelDensityOnGrid(gpu, st.models[0], leaders, 4);
        REQUIRE(ModelDensityMetalGrids() == before + 1);   // the GPU put it on the grid
        REQUIRE(gpu.grid.data.size() == cpu.grid.data.size());
        size_t different = 0;
        for (size_t i = 0; i < cpu.grid.data.size(); i++)
            different += std::memcmp(&cpu.grid.data[i], &gpu.grid.data[i], sizeof(float)) != 0;
        CHECK(different == 0);
    }
}

TEST_CASE("ModelDensityMetal_BitIdenticalToCPU", "[portable][Metal]") {
    // Triclinic, so that every term of the orthogonalization is there; then a group with symmetry to reduce.
    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