// 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 <algorithm>
#include <cmath>
#include <random>
#include <vector>

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

#include "../image_analysis/structure_refinement/ModelFFT.h"
#include "../image_analysis/structure_refinement/MolecularReplacement.h"
#include "../image_analysis/structure_refinement/MolecularReplacementMetal.h"

// The rotation function on the Apple GPU is the CPU's in float: equal to rounding, with the same best orientation.
// 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. The search as in MolecularReplacementTest.cpp.

namespace {
    // A synthetic body: carbon atoms scattered in a 12 A sphere, centroid at the origin.
    SearchModel synthetic_model() {
        SearchModel m;
        std::mt19937 rng(7);
        std::uniform_real_distribution<double> u(-12.0, 12.0);
        gemmi::Model model(1);
        gemmi::Chain chain("A");
        gemmi::Residue res;
        res.name = "UNK";
        res.seqid = gemmi::SeqId(1, ' ');
        std::vector<gemmi::Vec3> pos;
        while (pos.size() < 300) {
            const gemmi::Vec3 p(u(rng), u(rng), u(rng));
            if (p.length() < 12.0)
                pos.push_back(p);
        }
        gemmi::Vec3 c;
        for (const auto &p : pos)
            c += p;
        c *= 1.0 / pos.size();
        for (const auto &p : pos) {
            gemmi::Atom a;
            a.name = "CA";   // every atom counts in the clash check
            a.element = gemmi::El::C;
            a.pos = gemmi::Position(p - c);
            a.b_iso = 15.0f;
            a.occ = 1.0f;
            res.atoms.push_back(a);
            m.radius_A = std::max(m.radius_A, (p - c).length());
        }
        chain.residues.push_back(res);
        model.chains.push_back(chain);
        m.st.models.push_back(model);
        m.residues_read = m.residues_kept = 1;
        m.mass_da = 3600;
        return m;
    }

    gemmi::Mat33 axis_angle(gemmi::Vec3 axis, double deg) {
        axis = axis.normalized();
        const double k = deg * M_PI / 180.0, s = std::sin(k), c1 = 1.0 - std::cos(k);
        const double x = axis.x, y = axis.y, z = axis.z;
        return {1 + c1 * (x * x - 1), -s * z + c1 * x * y, s * y + c1 * x * z,
                s * z + c1 * x * y, 1 + c1 * (y * y - 1), -s * x + c1 * y * z,
                -s * y + c1 * x * z, s * x + c1 * y * z, 1 + c1 * (z * z - 1)};
    }

    // |F|^2 of the placed crystal to 4 A, as merged reflections.
    std::vector<MergedReflection> observed(const gemmi::Structure &crystal) {
        gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
        dc.d_min = 3.5;
        dc.rate = 1.5;
        dc.set_grid_cell_and_spacegroup(crystal);
        dc.put_model_density_on_grid(crystal.models[0]);
        const auto f = MapToFPhi(dc.grid).prepare_asu_data(3.5, 0, false, false, false);
        std::vector<MergedReflection> out;
        for (const auto &hv : f.v) {
            MergedReflection r;
            r.h = hv.hkl[0];
            r.k = hv.hkl[1];
            r.l = hv.hkl[2];
            r.I = std::norm(hv.value);
            r.sigma = 1;
            r.F = std::abs(hv.value);
            r.d = 1.0 / std::sqrt(crystal.cell.calculate_1_d2(hv.hkl));
            out.push_back(r);
        }
        return out;
    }
}

TEST_CASE("MolecularReplacementMetal_MatchesCPU", "[portable][MolecularReplacement][Metal]") {
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50);
    const auto merged = observed(PlaceSearchModel(model, truth, gemmi::Fractional(0.17, 0.31, 0.09), cell, sg));
    const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0);
    const MolecularTransform ft(model, 4.0, 4);
    REQUIRE(MolecularReplacementMetal::Available(ft));
    const MolecularReplacementMetal metal(ft, 4);

    const auto pg = CartesianPointGroup(sg, cell);
    const auto rotations = RotationAsu(UniformRotations(6.0), pg);
    const auto rf_cpu = RotationFunction(ft, refl, cell, sg, rotations, 4);
    const auto rf_metal = metal.RotationFunction(refl, cell, sg, rotations);
    REQUIRE(rf_metal.size() == rf_cpu.size());
    double worst = 0;
    for (size_t i = 0; i < rf_cpu.size(); i++)
        worst = std::max(worst, std::fabs(rf_metal[i] - rf_cpu[i]));
    CHECK(worst < 1e-4);
    const size_t best = std::max_element(rf_metal.begin(), rf_metal.end()) - rf_metal.begin();
    CHECK(best == static_cast<size_t>(std::max_element(rf_cpu.begin(), rf_cpu.end()) - rf_cpu.begin()));
    CHECK(RotationDistanceDeg(rotations[best], truth, pg) < 6.0);
}

#endif