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

// The rotation function's kernel (MolecularReplacementMetal.h), compiled when it is first used from this source,
// which CMakeLists.txt copies into the binary.

#include <metal_stdlib>
using namespace metal;

struct Params {
    int nu, nv, bu, bv;   // the grid, and its bricks along u and v
    uint nops, bins;
    float box, blur;
};

// MolecularTransform::At()'s corner, from the bricked grid.
static float2 transform_at(device const float2 *g, constant Params &p, int i, int j, int l) {
    bool conj = false;
    if (l < 0) {   // the half-l grid holds l >= 0; F(-h) = conj F(h)
        i = -i; j = -j; l = -l;
        conj = true;
    }
    if (i < 0 || i >= p.nu) { i %= p.nu; if (i < 0) i += p.nu; }
    if (j < 0 || j >= p.nv) { j %= p.nv; if (j < 0) j += p.nv; }
    const uint brick = uint(i >> 1) + uint(p.bu) * (uint(j >> 1) + uint(p.bv) * uint(l >> 1));
    float2 v = g[brick * 8 + uint(i & 1) + 2 * uint(j & 1) + 4 * uint(l & 1)];
    if (conj)
        v.y = -v.y;
    return v;
}

// One threadgroup per orientation, as MolecularReplacementGPU.cu's rotation_function; the reflections in the order
// of their bins, every bin's sums taken in a fixed order: each thread's own, then the SIMD group's, then the groups'.
// For orientation r, bin b: out[(r bins + b) 3 + k] = the sum over the bin's reflections of ic, ic^2 and
// ic (E^2 - 1 - mean), ic being |F|^2 summed over the copies.
kernel void rotation_function(device const float2 *g [[buffer(0)]], constant Params &p [[buffer(1)]],
                              device const float3x3 *rot [[buffer(2)]], device const packed_float3 *s [[buffer(3)]],
                              device const float *w [[buffer(4)]], device const uint *bin_first [[buffer(5)]],
                              device float *out [[buffer(6)]],
                              uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]],
                              uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]],
                              uint sg [[simdgroup_index_in_threadgroup]], uint nsg [[simdgroups_per_threadgroup]]) {
    threadgroup float sh[3][32];
    const float3x3 m = rot[r];   // R^T: the model turned by R has F_R(s) = F_model(R^T s)
    for (uint b = 0; b < p.bins; b++) {
        float s1 = 0, s2 = 0, sw = 0;
        for (uint i = bin_first[b] + t; i < bin_first[b + 1]; i += nt) {
            float ic = 0;
            for (uint j = 0; j < p.nops; j++) {
                const float3 v = m * float3(s[i * p.nops + j]);
                const float3 u = v * p.box;
                const float3 fl = floor(u);
                const float3 f = u - fl;
                const int3 i0 = int3(fl);
                float2 acc = 0;
                for (int dx = 0; dx < 2; dx++)
                    for (int dy = 0; dy < 2; dy++)
                        for (int dz = 0; dz < 2; dz++) {
                            const float wt = (dx ? f.x : 1 - f.x) * (dy ? f.y : 1 - f.y) * (dz ? f.z : 1 - f.z);
                            acc += wt * transform_at(g, p, i0.x + dx, i0.y + dy, i0.z + dz);
                        }
                // |F exp(blur s^2 / 4)|^2
                ic += dot(acc, acc) * exp(p.blur * dot(v, v) / 2.0f);
            }
            s1 += ic;
            s2 += ic * ic;
            sw += ic * w[i];
        }
        s1 = simd_sum(s1);
        s2 = simd_sum(s2);
        sw = simd_sum(sw);
        if (lane == 0) {
            sh[0][sg] = s1;
            sh[1][sg] = s2;
            sh[2][sg] = sw;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        if (t == 0) {
            float a = 0, c = 0, d = 0;
            for (uint k = 0; k < nsg; k++) {
                a += sh[0][k];
                c += sh[1][k];
                d += sh[2][k];
            }
            device float *o = out + (r * p.bins + b) * 3;
            o[0] = a;
            o[1] = c;
            o[2] = d;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
}