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

// The model density's gather (ModelDensityMetal.h), compiled when it is first used from this source, which
// CMakeLists.txt copies into the binary - with fast math off and no a*b+c contracted, as on the host, so that
// every operation below rounds as the CPU's does.

#include <metal_stdlib>
using namespace metal;

#pragma clang fp contract(off)

// A brick's edge in points (ModelDensityMetal.mm says the same): every point of a brick tests each atom whose box
// reaches the brick, and with 4 rather than CUDA's 8 far fewer of them test atoms that miss them - the kernel
// took a half to a third of the time on an Apple M6.
constant int BRICK = 4;

struct Atom {   // ModelDensityMetalAtom
    float nctr[3];
    int c[3], d[3];
    float max_r2;
    float occ;
    int aniso;
    float a[5];
    float b[5][6];
};

// Where an atom's box starts along each axis, wrapped into the grid, and its width less one (2 d): the test
// whether a point is in the box, without a division.
struct Box {
    int first[3];
    int span[3];
};

struct Grid {   // ModelDensityMetalGrid
    int nu, nv, nw;
    float m00, m01, m02, m11, m12, m22;
};

// gemmi's unsafe_expapprox() (formfact.hpp, MPL-2.0; https://github.com/project-gemmi/gemmi), which is based on
// expapprox() from https://github.com/jhjourdan/SIMD-math-prims: the same operations, the bits moved with as_type
// where gemmi uses memcpy.
static float expapprox(float x) {
    const float val = 12102203.1615614f * x + 1065353216.f;
    const int vali = static_cast<int>(val);
    const int xu1 = vali & 0x7F800000;
    const int xu2 = (vali & 0x7FFFFF) | 0x3F800000;
    const float a = as_type<float>(xu1);
    const float b = as_type<float>(xu2);
    return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b *
            (-2.190619930e-3f + b * 1.3555747234e-2f))));
}

// The point p's offset in the box along an axis of n points, 0 to span where it is in it.
static int box_offset(int p, int first, int n) {
    const int m = p - first;
    return m < 0 ? m + n : m;
}

// One threadgroup per brick of BRICK^3 points: each point adds the brick's atoms in model order, as
// UseDensityBoxesByPlane() and AtomDensity (ModelGrid.cpp) add them on the CPU.
kernel void model_density(device const Atom *atoms [[buffer(0)]], constant Grid &g [[buffer(1)]],
                          device const uint *pair_atom [[buffer(2)]], device const uint *brick_first [[buffer(3)]],
                          device float *out [[buffer(4)]], device const Box *boxes [[buffer(5)]],
                          uint3 group [[threadgroup_position_in_grid]], uint3 t [[thread_position_in_threadgroup]]) {
    const uint brick = group.x;
    const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK;
    const int u = int(brick % uint(nbu)) * BRICK + int(t.x);
    const int v = int(brick / uint(nbu) % uint(nbv)) * BRICK + int(t.y);
    const int w = int(brick / uint(nbu * nbv)) * BRICK + int(t.z);
    if (u >= g.nu || v >= g.nv || w >= g.nw)
        return;
    float acc = 0.0f;
    for (uint k = brick_first[brick]; k < brick_first[brick + 1]; k++) {
        const uint i = pair_atom[k];
        device const Box &box = boxes[i];
        const int mu = box_offset(u, box.first[0], g.nu), mv = box_offset(v, box.first[1], g.nv),
                  mw = box_offset(w, box.first[2], g.nw);
        if (mu > box.span[0] || mv > box.span[1] || mw > box.span[2])
            continue;
        device const Atom &at = atoms[i];
        // The point's index in the box's own (unwrapped) range, as the CPU walks it.
        const int uu = at.c[0] - at.d[0] + mu, vv = at.c[1] - at.d[1] + mv, ww = at.c[2] - at.d[2] + mw;
        const float fx = at.nctr[0] - float(uu), fy = at.nctr[1] - float(vv), fz = at.nctr[2] - float(ww);
        const float z = g.m22 * fz;
        const float y = g.m11 * fy + g.m12 * fz;
        const float yz = y * y + z * z;
        const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
        const float r2 = yz + x * x;
        if (r2 > at.max_r2)
            continue;
        float density = 0.0f;
        if (at.aniso == 0) {
            for (int q = 0; q < 5; q++)
                density += at.a[q] * expapprox(max(at.b[q][0] * r2, -88.f));
        } else {
            for (int q = 0; q < 5; q++) {
                device const float *b = at.b[q];
                // gemmi's SMat33::r_u_r(), in its order
                const float rur = x * x * b[0] + y * y * b[1] + z * z * b[2]
                                  + 2 * (x * y * b[3] + x * z * b[4] + y * z * b[5]);
                density += at.a[q] * expapprox(max(rur, -88.f));
            }
        }
        acc += at.occ * density;
    }
    out[uint(u) + uint(g.nu) * (uint(v) + uint(g.nv) * uint(w))] = acc;
}