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

#include "ModelDensityGPU.h"

#include <algorithm>
#include <climits>
#include <cmath>

#include <cub/cub.cuh>

#include "../../common/JFJochException.h"

namespace {

void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

constexpr int BRICK = 8;                  // the gather's bricks are BRICK^3 grid points, one block each
constexpr int GATHER_TILE = 128;           // atoms staged in shared memory at a time
constexpr int MAX_BRICKS_PER_AXIS = 16;   // an atom's box touches at most this many bricks along an axis
constexpr int THREADS = 256;
constexpr int SORT_ATOMS = 4096;          // a brick's atoms are sorted in shared memory up to this many

__device__ __forceinline__ int imod(int a, int n) {
    const int r = a % n;
    return r < 0 ? r + n : r;
}

// gemmi's unsafe_expapprox() (formfact.hpp), the exponential the density is computed with.
__device__ __forceinline__ float expapprox(float x) {
    const float val = 12102203.1615614f * x + 1065353216.f;
    const int vali = static_cast<int>(val);
    const float a = __int_as_float(vali & 0x7F800000);
    const float b = __int_as_float((vali & 0x7FFFFF) | 0x3F800000);
    return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * (-2.190619930e-3f + b * 1.3555747234e-2f))));
}

// The distinct bricks the points c - d ... c + d of one axis fall in, wrapped into the cell. A grid size
// that is not a multiple of BRICK leaves the last brick partial, which is why this is done on wrapped
// points and not in brick coordinates. -1 where there are more than MAX_BRICKS_PER_AXIS of them.
__device__ int axis_bricks(int c, int d, int n, int *out) {
    int k = 0;
    for (int p = c - d; p <= c + d; p++) {
        const int b = imod(p, n) / BRICK;
        bool seen = false;
        for (int j = 0; j < k; j++)
            seen = seen || out[j] == b;
        if (seen)
            continue;
        if (k == MAX_BRICKS_PER_AXIS)
            return -1;
        out[k++] = b;
    }
    return k;
}

// The bricks atom i's box touches, per axis. The box is gemmi's: `box` points either side of the
// point nearest the atom.
__device__ void atom_bricks(const float4 &pos, const int *box, int nu, int nv, int nw, int *bu, int &nbu, int *bv,
                            int &nbv, int *bw, int &nbw) {
    nbu = axis_bricks(__float2int_rn(pos.x), box[0], nu, bu);
    nbv = axis_bricks(__float2int_rn(pos.y), box[1], nv, bv);
    nbw = axis_bricks(__float2int_rn(pos.z), box[2], nw, bw);
}

// How many atoms touch each brick: integer atomics, whose sum does not depend on their order.
__global__ void count_bricks_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *brick_count) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c;
    atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c);
    const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK;
    for (int k = 0; k < c; k++)
        for (int j = 0; j < b; j++)
            for (int l = 0; l < a; l++)
                atomicAdd(brick_count + bu[l] + nbu * (bv[j] + nbv * bw[k]), 1);
}

// Each atom into the segments of the bricks it touches, at whatever place in the segment the atomic hands out:
// the order within a brick is settled by order_bricks_kernel.
__global__ void fill_bricks_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw,
                                   const int *brick_start, int *brick_fill, int *pair_atom) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c;
    atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c);
    const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK;
    for (int k = 0; k < c; k++)
        for (int j = 0; j < b; j++)
            for (int l = 0; l < a; l++) {
                const int brick = bu[l] + nbu * (bv[j] + nbv * bw[k]);
                pair_atom[brick_start[brick] + atomicAdd(brick_fill + brick, 1)] = i;
            }
}

// Each brick's atoms into model order, one block per brick: a bitonic sort of the segment in shared memory,
// padded to a power of two, for the bricks of up to SORT_ATOMS atoms. A larger brick is ranked instead: an
// atom's place is the number of the brick's atoms with a lower index (an atom is in a brick once), counted
// against the segment staged a tile at a time, a pair staying with one thread through the tiles.
__global__ void order_bricks_kernel(const int *unordered, const int *brick_start, int *rank, int *pair_atom) {
    const int brick = blockIdx.x;
    const int s = brick_start[brick], e = brick_start[brick + 1], k = e - s;
    __shared__ int keys[SORT_ATOMS];
    if (k <= SORT_ATOMS) {
        int n = 1;
        while (n < k)
            n <<= 1;
        for (int t = threadIdx.x; t < n; t += THREADS)
            keys[t] = t < k ? unordered[s + t] : INT_MAX;
        __syncthreads();
        for (int size = 2; size <= n; size <<= 1)
            for (int stride = size >> 1; stride > 0; stride >>= 1) {
                for (int t = threadIdx.x; t < n / 2; t += THREADS) {
                    const int i = 2 * t - (t & (stride - 1)), j = i + stride;
                    const int a = keys[i], b = keys[j];
                    if ((a > b) == ((i & size) == 0)) {
                        keys[i] = b;
                        keys[j] = a;
                    }
                }
                __syncthreads();
            }
        for (int t = threadIdx.x; t < k; t += THREADS)
            pair_atom[s + t] = keys[t];
        return;
    }
    for (int c = s; c < e; c += THREADS) {
        const int m = min(THREADS, e - c);
        __syncthreads();
        if (threadIdx.x < m)
            keys[threadIdx.x] = unordered[c + threadIdx.x];
        __syncthreads();
        for (int p = s + threadIdx.x; p < e; p += THREADS) {
            const int atom = unordered[p];
            int below = 0;
            for (int q = 0; q < m; q++)
                below += keys[q] < atom;
            rank[p] = (c == s ? 0 : rank[p]) + below;
        }
    }
    for (int p = s + threadIdx.x; p < e; p += THREADS)
        pair_atom[s + rank[p]] = unordered[p];
}

// What the gather needs of one staged atom: its isotropic Gaussians (an anisotropic atom's matrices are read
// from the atoms in global memory), and its box as the brick sees it - a bit per point of the brick along each
// axis for whether the point is in the box, and for the points that are, what gemmi's box offset makes of them:
// the point's unwrapped box index is base + t, less n where t reaches wrap.
struct GatherAtom {
    float a[5], b[5];
    float occ, max_r2;
    int index, aniso;
    int masks;         // bits 0-7: u, 8-15: v, 16-23: w
    int base[3], wrap[3];
    float nctr[3];
};

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

// The density of the model on the grid, as UseDensityBoxesByPlane() and AtomDensity (ModelGrid.cpp) put
// it there, bit for bit: every point adds, in model order, each atom whose box holds it and whose radius
// reaches it, with the CPU's single-precision arithmetic in the CPU's order (no a*b+c is contracted:
// CMakeLists.txt compiles device code with --fmad=false). One block per brick; the brick's atoms are
// staged in shared memory. Supports() leaves only boxes narrower than the cell, so a point is in an atom's
// box at most once, at one index of the box's own unwrapped range.
__global__ void gather_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *box, const int *pair_atom,
                              const int *brick_start, const int *brick_end, ModelDensityGPU::Geometry g,
                              float *grid) {
    const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK;
    const int brick = blockIdx.x;
    const int s = brick_start[brick], e = brick_end[brick];   // s == e: an empty brick, written as zeros
    const int origin[3] = {(brick % nbu) * BRICK, (brick / nbu % nbv) * BRICK, (brick / (nbu * nbv)) * BRICK};
    const int tu = threadIdx.x, tv = threadIdx.y, tw = threadIdx.z;
    const int u = origin[0] + tu, v = origin[1] + tv, w = origin[2] + tw;
    const bool inside = u < g.nu && v < g.nv && w < g.nw;
    const int tid = tu + BRICK * (tv + BRICK * tw);
    const int n3[3] = {g.nu, g.nv, g.nw};

    __shared__ GatherAtom sh[GATHER_TILE];
    float acc = 0.0f;
    for (int c = s; c < e; c += GATHER_TILE) {
        const int m = min(GATHER_TILE, e - c);
        __syncthreads();
        if (tid < m) {
            const int a = pair_atom[c + tid];
            const ModelDensityAtom &at = atoms[a];
            const float4 p = pos[a];
            const float nctr[3] = {p.x, p.y, p.z};
            GatherAtom &r = sh[tid];
            for (int q = 0; q < 5; q++) {
                r.a[q] = at.a[q];
                r.b[q] = at.b[q][0];
            }
            r.occ = at.occ;
            r.max_r2 = at.max_r2;
            r.index = a;
            r.aniso = at.aniso;
            r.masks = 0;
            for (int k = 0; k < 3; k++) {
                const int d = box[3 * a + k];
                const int lo = __float2int_rn(nctr[k]) - d, first = imod(lo, n3[k]), span = 2 * d;
                const int m0 = box_offset(origin[k], first, n3[k]);
                for (int t = 0; t < BRICK; t++) {
                    const int mt = m0 + t;
                    if ((mt < n3[k] ? mt : mt - n3[k]) <= span)
                        r.masks |= 1 << (8 * k + t);
                }
                r.base[k] = lo + m0;
                r.wrap[k] = n3[k] - m0;
                r.nctr[k] = nctr[k];
            }
        }
        __syncthreads();
        if (!inside)
            continue;
        for (int k = 0; k < m; k++) {
            const GatherAtom &at = sh[k];
            if (!((at.masks >> tu) & (at.masks >> (8 + tv)) & (at.masks >> (16 + tw)) & 1))
                continue;
            const int pu = at.base[0] + tu - (tu >= at.wrap[0] ? g.nu : 0);
            const int pv = at.base[1] + tv - (tv >= at.wrap[1] ? g.nv : 0);
            const int pw = at.base[2] + tw - (tw >= at.wrap[2] ? g.nw : 0);
            // f = centre - (u, v, w) in the box's own (unwrapped) indices, then in A, as the CPU walks it
            const float fx = at.nctr[0] - static_cast<float>(pu);
            const float fy = at.nctr[1] - static_cast<float>(pv);
            const float fz = at.nctr[2] - static_cast<float>(pw);
            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) {
                for (int q = 0; q < 5; q++)
                    density += at.a[q] * expapprox(fmaxf(at.b[q] * r2, -88.f));
            } else {
                for (int q = 0; q < 5; q++) {
                    const float *bq = atoms[at.index].b[q];
                    // gemmi's SMat33::r_u_r(), in its order
                    const float rur = x * x * bq[0] + y * y * bq[1] + z * z * bq[2] +
                                      2 * (x * y * bq[3] + x * z * bq[4] + y * z * bq[5]);
                    density += at.a[q] * expapprox(fmaxf(rur, -88.f));
                }
            }
            acc += at.occ * density;
        }
    }
    if (inside)
        grid[u + static_cast<size_t>(g.nu) * (v + static_cast<size_t>(g.nv) * w)] = acc;
}

int blocks(size_t n) {
    return static_cast<int>((n + THREADS - 1) / THREADS);
}

// cub's temporary storage for the scan over the bricks.
size_t CubBytes(size_t bricks) {
    size_t scan = 0;
    cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
                                           static_cast<int>(bricks + 1)));
    return scan;
}

}  // namespace

// The most bricks the 2 d + 1 wrapped points of a box can fall in along an axis of n points: one more
// than the points span for where they start in a brick, and one more again where they wrap past a last
// brick that is partial (n = 17: the points 15, 16, 0 fall in bricks 1, 2 and 0).
size_t ModelDensityGPU::AxisBrickBound(int d, int n) {
    const int nb = (n + BRICK - 1) / BRICK;
    return std::min(nb, (2 * d + 1 + BRICK - 1) / BRICK + 2);
}

size_t ModelDensityGPU::Bricks(int nu, int nv, int nw) {
    return static_cast<size_t>((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK);
}

void ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const ModelDensityAtom *atoms, size_t n, int *box) {
    const int n3[3] = {grid.nu, grid.nv, grid.nw};
    double spacing[3];
    for (int k = 0; k < 3; k++)
        spacing[k] = 1.0 / (n3[k] * std::sqrt(grid.frac[3 * k] * grid.frac[3 * k] + grid.frac[3 * k + 1] * grid.frac[3 * k + 1] +
                                              grid.frac[3 * k + 2] * grid.frac[3 * k + 2]));
    for (size_t i = 0; i < n; i++)
        for (int k = 0; k < 3; k++)
            box[3 * i + k] = static_cast<int>(std::ceil(atoms[i].radius / spacing[k]));
}

std::vector<int> ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
    std::vector<int> box(3 * atoms.size());
    Boxes(grid, atoms.data(), atoms.size(), box.data());
    return box;
}

size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const int *box, size_t n) {
    const int n3[3] = {grid.nu, grid.nv, grid.nw};
    size_t pairs = 0;
    for (size_t i = 0; i < n; i++) {
        size_t product = 1;
        for (int k = 0; k < 3; k++)
            product *= AxisBrickBound(box[3 * i + k], n3[k]);
        pairs += product;
    }
    return pairs;
}

size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
    const std::vector<int> box = Boxes(grid, atoms);
    return PairBound(grid, box.data(), atoms.size());
}

bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const int *box, size_t n) {
    const int n3[3] = {grid.nu, grid.nv, grid.nw};
    for (size_t i = 0; i < 3 * n; i++)
        if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS)
            return false;
    return true;
}

bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
    const std::vector<int> box = Boxes(grid, atoms);
    return Supports(grid, box.data(), atoms.size());
}

size_t ModelDensityGPU::DeviceBytes(size_t max_atoms, size_t max_pairs, size_t max_bricks) {
    const size_t na = max_atoms + 1, nb = max_bricks + 1;
    return na * (sizeof(ModelDensityAtom) + 3 * sizeof(int)) + 3 * max_pairs * sizeof(int) + 3 * nb * sizeof(int) +
           CubBytes(nb);
}

ModelDensityGPU::ModelDensityGPU(cudaStream_t stream, size_t max_atoms_, size_t max_pairs_, size_t max_bricks_)
    : stream(stream), max_atoms(max_atoms_), max_pairs(max_pairs_), max_bricks(max_bricks_) {
    const auto sync = CudaAlloc::Synchronous;
    const size_t na = std::max<size_t>(max_atoms, 1), np = std::max<size_t>(max_pairs, 1), nb = max_bricks + 1;
    atoms = CudaDevicePtr<ModelDensityAtom>(na, sync);
    box = CudaDevicePtr<int>(3 * na, sync);
    brick_count = CudaDevicePtr<int>(nb, sync);
    brick_start = CudaDevicePtr<int>(nb, sync);
    brick_fill = CudaDevicePtr<int>(nb, sync);
    unordered = CudaDevicePtr<int>(np, sync);
    rank = CudaDevicePtr<int>(np, sync);
    pair_atom = CudaDevicePtr<int>(np, sync);
    cub_bytes = CubBytes(nb);
    cub_temp = CudaDevicePtr<char>(std::max<size_t>(cub_bytes, 1), sync);
    host_box = CudaHostPtr<int>(3 * na);
}

void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atom_list) {
    SetAtoms(grid, atom_list.data(), atom_list.size());
}

void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const ModelDensityAtom *atom_list, size_t n) {
    if (n > max_atoms || Bricks(grid.nu, grid.nv, grid.nw) > max_bricks)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: grid or atoms over the reserve");
    // The boxes into the pinned staging buffer, so that their copy and the atoms' (pinned where the caller keeps
    // them so) go up by DMA instead of through the driver's staging of pageable memory.
    Boxes(grid, atom_list, n, host_box);
    // An atom whose box is wider than the cell would reach a point through more than one image, which
    // gemmi's box walk adds and the gather's box test does not; such a cell is left to the CPU.
    if (!Supports(grid, host_box, n))
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom wider than the cell");
    if (PairBound(grid, host_box, n) > max_pairs)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: gather pairs over the reserve");
    geom.nu = grid.nu;
    geom.nv = grid.nv;
    geom.nw = grid.nw;
    // As gemmi's Grid sets orth_n (multiply_by_diagonal with 1/n) and ModelDensityGeometry() rounds it.
    const double inv_n[3] = {1.0 / grid.nu, 1.0 / grid.nv, 1.0 / grid.nw};
    geom.m00 = static_cast<float>(grid.orth[0] * inv_n[0]);
    geom.m01 = static_cast<float>(grid.orth[1] * inv_n[1]);
    geom.m02 = static_cast<float>(grid.orth[2] * inv_n[2]);
    geom.m11 = static_cast<float>(grid.orth[4] * inv_n[1]);
    geom.m12 = static_cast<float>(grid.orth[5] * inv_n[2]);
    geom.m22 = static_cast<float>(grid.orth[8] * inv_n[2]);
    n_atoms = static_cast<int>(n);
    if (n_atoms > 0) {
        cuda_err(cudaMemcpyAsync(atoms, atom_list, n * sizeof(ModelDensityAtom), cudaMemcpyHostToDevice, stream));
        cuda_err(cudaMemcpyAsync(box, host_box, 3 * n * sizeof(int), cudaMemcpyHostToDevice, stream));
    }
    // The uploads read the caller's atoms, which it may change once this returns.
    cuda_err(cudaStreamSynchronize(stream));
}

void ModelDensityGPU::Compute(const float4 *d_pos, float *d_grid) {
    // Counted per brick, scanned, filled, and each brick put in model order: nothing comes back to the host, so
    // the launches queue straight through (and record into a CUDA graph). The bricks' atoms end up as a stable
    // sort of the atom-major pairs by brick would leave them.
    const int nb = static_cast<int>(Bricks(geom.nu, geom.nv, geom.nw));
    cuda_err(cudaMemsetAsync(brick_count, 0, (nb + 1) * sizeof(int), stream));
    cuda_err(cudaMemsetAsync(brick_fill, 0, nb * sizeof(int), stream));
    if (n_atoms > 0) {
        count_bricks_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
                                                                     brick_count);
        cuda_err(cudaGetLastError());
    }
    size_t temp = cub_bytes;
    cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, brick_count.get(), brick_start.get(), nb + 1, stream));
    if (n_atoms > 0) {
        fill_bricks_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
                                                                    brick_start, brick_fill, unordered);
        cuda_err(cudaGetLastError());
    }
    order_bricks_kernel<<<nb, THREADS, 0, stream>>>(unordered, brick_start, rank, pair_atom);
    cuda_err(cudaGetLastError());
    gather_kernel<<<nb, dim3(BRICK, BRICK, BRICK), 0, stream>>>(atoms, d_pos, box, pair_atom, brick_start,
                                                                brick_start.get() + 1, geom, d_grid);
    cuda_err(cudaGetLastError());
}