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

// Sparse connected-component labelling adapted from the design the ACTS/traccc project arrived at
// for sparse silicon-detector clusterization (backward-neighbour graph over a sorted hit list, then
// a parallel union-find), which is itself the GPU counterpart of the SparseCCL that
// StrongPixelSet.cpp uses on the host.

#include <climits>

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

namespace {

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

constexpr int THREADS = 256;
constexpr int FINISH_THREADS = 1024;

// --- compaction: packed bit buffer -> strong-pixel list, sorted by flat index -------------------
// Each block owns a CONTIGUOUS range of words. Pass 1 counts its strong bits, a single-block scan
// turns the counts into offsets, and pass 2 walks the same range in increasing order and writes at
// that offset. No atomics anywhere, which is what keeps the output sorted.
__global__ void count_bits(const uint32_t *__restrict__ strong, const uint32_t *__restrict__ res_mask,
                           uint32_t *__restrict__ block_count, size_t nwords) {
    const size_t per_block = (nwords + gridDim.x - 1) / gridDim.x;
    const size_t w0 = static_cast<size_t>(blockIdx.x) * per_block;
    const size_t w1 = min(w0 + per_block, nwords);
    uint32_t local = 0;
    for (size_t w = w0 + threadIdx.x; w < w1; w += blockDim.x)
        local += __popc(strong[w] & ~res_mask[w]);

    __shared__ uint32_t block_total;
    if (threadIdx.x == 0) block_total = 0;
    __syncthreads();
    atomicAdd(&block_total, local);
    __syncthreads();
    if (threadIdx.x == 0) block_count[blockIdx.x] = block_total;
}

__global__ void scan_block_counts(const uint32_t *__restrict__ in, uint32_t *__restrict__ out,
                                  uint32_t *__restrict__ total, int n) {
    __shared__ uint32_t shared[FINISH_THREADS];
    const int t = threadIdx.x, nthreads = blockDim.x;
    const int chunk = (n + nthreads - 1) / nthreads;
    const int lo = min(t * chunk, n), hi = min(lo + chunk, n);
    uint32_t sum = 0;
    for (int i = lo; i < hi; i++) sum += in[i];
    shared[t] = sum;
    __syncthreads();
    for (int d = 1; d < nthreads; d <<= 1) {
        const uint32_t v = (t >= d) ? shared[t - d] : 0u;
        __syncthreads();
        shared[t] += v;
        __syncthreads();
    }
    uint32_t acc = shared[t] - sum;
    for (int i = lo; i < hi; i++) { out[i] = acc; acc += in[i]; }
    if (t == nthreads - 1) *total = shared[t];
}

// One thread per block emits its range. Strong pixels are ~1e-4 of the image, so a block's range
// holds a handful of them and serial emission is both trivially ordered and fast; the parallelism
// comes from the block count.
__global__ void scatter_bits(const uint32_t *__restrict__ strong, const uint32_t *__restrict__ res_mask,
                             const uint32_t *__restrict__ block_offset, const int32_t *__restrict__ image,
                             uint32_t *__restrict__ out_index, int32_t *__restrict__ out_value,
                             size_t nwords, uint32_t capacity) {
    if (threadIdx.x != 0) return;
    const size_t per_block = (nwords + gridDim.x - 1) / gridDim.x;
    const size_t w0 = static_cast<size_t>(blockIdx.x) * per_block;
    const size_t w1 = min(w0 + per_block, nwords);
    uint32_t pos = block_offset[blockIdx.x];
    for (size_t w = w0; w < w1; ++w) {
        uint32_t word = strong[w] & ~res_mask[w];
        while (word != 0) {
            const uint32_t flat = static_cast<uint32_t>(w * 32 + (__ffs(word) - 1));
            word &= word - 1;
            if (pos < capacity) {
                out_index[pos] = flat;
                out_value[pos] = image[flat];
            }
            ++pos;
        }
    }
}

// --- connected components ----------------------------------------------------------------------
__device__ __forceinline__ int lower_bound_device(const uint32_t *a, int n, uint32_t key) {
    int lo = 0, hi = n;
    while (lo < hi) {
        const int mid = (lo + hi) >> 1;
        if (a[mid] < key) lo = mid + 1; else hi = mid;
    }
    return lo;
}

// Walk to the root, halving the path on the way. The plain store is safe: a parent only ever
// decreases and the grandparent is always still an ancestor, so a concurrent writer can only make
// the chain shorter. Without halving, a long thin feature - a diffraction ring is exactly one -
// builds a parent chain as long as the feature itself and every later merge walks all of it.
__device__ __forceinline__ uint32_t find_root(uint32_t *parent, uint32_t a) {
    uint32_t p = parent[a];
    while (p != a) {
        const uint32_t gp = parent[p];
        if (gp == p) return p;
        parent[a] = gp;
        a = gp;
        p = parent[a];
    }
    return a;
}

// Lock-free union. Terminates because max(a,b) strictly decreases; converges on the component's
// LOWEST index as its root, which is what the host's make_union does too.
__device__ __forceinline__ void merge_roots(uint32_t *parent, uint32_t a, uint32_t b) {
    a = find_root(parent, a);
    b = find_root(parent, b);
    while (a != b) {
        if (a < b) { const uint32_t t = a; a = b; b = t; }
        const uint32_t old = atomicMin(&parent[a], b);
        if (old == a) return;                    // a was a root and now points at b: done
        a = find_root(parent, old);              // someone re-parented a; carry on from its root
        b = find_root(parent, b);
    }
}

__global__ void init_parent(uint32_t *__restrict__ parent, const uint32_t *__restrict__ nstrong,
                            uint32_t capacity) {
    const int n = static_cast<int>(min(*nstrong, capacity));
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
        parent[i] = i;
}

// The list is sorted by flat index, so a pixel's 8-neighbours that come EARLIER in it are exactly
// four: (line, col-1), (line-1, col-1), (line-1, col) and (line-1, col+1). Each is one binary
// search away, which is what makes the sparse formulation cheap.
__global__ void union_neighbours(const uint32_t *__restrict__ index, uint32_t *__restrict__ parent,
                                 const uint32_t *__restrict__ nstrong, uint32_t capacity, int width) {
    const int n = static_cast<int>(min(*nstrong, capacity));
    if (static_cast<uint32_t>(n) >= capacity) return;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        const uint32_t flat = index[i];
        const uint32_t col = flat % width;
        const uint32_t line = flat / width;
        uint32_t candidate[4];
        int ncandidate = 0;
        if (col > 0)                                      candidate[ncandidate++] = flat - 1;
        if (line > 0 && col > 0)                          candidate[ncandidate++] = flat - width - 1;
        if (line > 0)                                     candidate[ncandidate++] = flat - width;
        if (line > 0 && col + 1 < static_cast<uint32_t>(width)) candidate[ncandidate++] = flat - width + 1;
        for (int k = 0; k < ncandidate; k++) {
            const int j = lower_bound_device(index, i, candidate[k]);
            if (j < i && index[j] == candidate[k]) merge_roots(parent, static_cast<uint32_t>(i), static_cast<uint32_t>(j));
        }
    }
}

__global__ void resolve_roots(uint32_t *__restrict__ parent, uint32_t *__restrict__ root,
                              const uint32_t *__restrict__ nstrong, uint32_t capacity) {
    const int n = static_cast<int>(min(*nstrong, capacity));
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
        root[i] = find_root(parent, static_cast<uint32_t>(i));
}

// Everything after the labelling in ONE block, so a frame needs a single host synchronisation:
// hand out labels, count the members of each component, sum the surviving ones, filter by max-pix
// and compact - all of it order-preserving.
__global__ void finish_components(const uint32_t *__restrict__ index, const int32_t *__restrict__ value,
                                  const uint32_t *__restrict__ root, uint32_t *__restrict__ label,
                                  int32_t *__restrict__ count, SpotExtractorGPUSpot *__restrict__ scratch,
                                  SpotExtractorGPUSpot *__restrict__ out, uint32_t *__restrict__ nout,
                                  const uint32_t *__restrict__ nstrong, uint32_t capacity,
                                  int width, int max_pix) {
    const int n = static_cast<int>(min(*nstrong, capacity));
    if (threadIdx.x == 0) *nout = 0;
    __syncthreads();
    // Same give-up as StrongPixelSet::FindComponentsImage - except that here the count is known
    // before a single pixel has been written anywhere, so the frame costs nothing to reject.
    if (n == 0 || static_cast<uint32_t>(n) >= capacity) return;

    __shared__ uint32_t shared[FINISH_THREADS];
    const int t = threadIdx.x, nthreads = blockDim.x;
    const int chunk = (n + nthreads - 1) / nthreads;
    const int lo = min(t * chunk, n), hi = min(lo + chunk, n);

    // 1) labels, by a prefix sum over the roots in ascending order - the order the host's second
    //    scan hands them out in, which is what makes the spot ORDER identical.
    uint32_t nroot = 0;
    for (int i = lo; i < hi; i++) nroot += (root[i] == static_cast<uint32_t>(i)) ? 1u : 0u;
    shared[t] = nroot;
    __syncthreads();
    for (int d = 1; d < nthreads; d <<= 1) {
        const uint32_t v = (t >= d) ? shared[t - d] : 0u;
        __syncthreads();
        shared[t] += v;
        __syncthreads();
    }
    uint32_t next_label = shared[t] - nroot;
    for (int i = lo; i < hi; i++)
        if (root[i] == static_cast<uint32_t>(i)) label[i] = next_label++;
    const int nlabel = static_cast<int>(shared[nthreads - 1]);
    __syncthreads();

    // 2) member counts
    for (int i = t; i < nlabel; i += nthreads) count[i] = 0;
    __syncthreads();
    for (int i = t; i < n; i += nthreads) atomicAdd(&count[label[root[i]]], 1);
    __syncthreads();

    // 3) sums, one thread per component, walking its members in ascending list order so the
    //    accumulation matches DiffractionSpot::AddPixel term for term. A component bigger than
    //    max-pix is thrown away below, so it is not summed - which is also what keeps a whole lit
    //    module or diffraction ring from turning into one thread walking tens of thousands of
    //    entries.
    for (int i = t; i < n; i += nthreads) {
        if (root[i] != static_cast<uint32_t>(i)) continue;
        const uint32_t l = label[i];
        const int want = count[l];
        scratch[l].pixel_count = want;
        if (want > max_pix) continue;
        long long x = 0, y = 0;
        long long photons = 0, max_photons = LLONG_MIN;
        int found = 0;
        for (int j = i; j < n && found < want; j++) {
            if (root[j] != static_cast<uint32_t>(i)) continue;
            const long long counts = value[j];
            // Integers, exactly as DiffractionSpot::AddPixel does them, so host and device agree by
            // construction - no rounding mode to match and nothing for either compiler to contract.
            x += static_cast<long long>(index[j] % width) * counts;
            y += static_cast<long long>(index[j] / width) * counts;
            photons += counts;
            max_photons = max(max_photons, counts);
            found++;
        }
        scratch[l].x = x;
        scratch[l].y = y;
        scratch[l].photons = photons;
        scratch[l].max_photons = max_photons;
    }
    __syncthreads();

    // 4) max-pix filter, compacted by another prefix sum so the surviving spots keep their order
    const int label_chunk = (nlabel + nthreads - 1) / nthreads;
    const int label_lo = min(t * label_chunk, nlabel), label_hi = min(label_lo + label_chunk, nlabel);
    uint32_t nkeep = 0;
    for (int i = label_lo; i < label_hi; i++) nkeep += (scratch[i].pixel_count <= max_pix) ? 1u : 0u;
    shared[t] = nkeep;
    __syncthreads();
    for (int d = 1; d < nthreads; d <<= 1) {
        const uint32_t v = (t >= d) ? shared[t - d] : 0u;
        __syncthreads();
        shared[t] += v;
        __syncthreads();
    }
    uint32_t pos = shared[t] - nkeep;
    for (int i = label_lo; i < label_hi; i++)
        if (scratch[i].pixel_count <= max_pix) out[pos++] = scratch[i];
    if (t == nthreads - 1) *nout = shared[t];
}

} // namespace

SpotExtractorGPU::SpotExtractorGPU(int32_t in_width, int32_t in_height, std::shared_ptr<CudaStream> in_stream)
    : stream(std::move(in_stream)),
      width(in_width),
      nwords((static_cast<size_t>(in_width) * in_height + 31) / 32),
      gpu_res_mask(nwords),
      gpu_nstrong(1),
      gpu_index(MAX_STRONG),
      gpu_value(MAX_STRONG),
      gpu_parent(MAX_STRONG),
      gpu_root(MAX_STRONG),
      gpu_label(MAX_STRONG),
      gpu_count(MAX_STRONG),
      gpu_spot(MAX_STRONG),
      gpu_spot_out(MAX_STRONG),
      gpu_nspot(1),
      host_nspot(1),
      host_spot(SPOT_PREFIX) {
    // One block per few hundred words: enough blocks to fill the device, few enough that the serial
    // emission inside a block stays short even when a whole detector row lights up.
    compact_blocks = static_cast<int>((nwords + 255) / 256);
    if (compact_blocks > 4096) compact_blocks = 4096;
    if (compact_blocks < 1) compact_blocks = 1;
    gpu_block_count = CudaDevicePtr<uint32_t>(compact_blocks);
    gpu_block_offset = CudaDevicePtr<uint32_t>(compact_blocks);

    // Nothing excluded except the padding bits of the last word - the same starting point as
    // ImageSpotFinder's own mask, so the two agree even when no resolution mask is ever set.
    std::vector<uint32_t> mask(nwords, 0);
    const size_t npixel = static_cast<size_t>(in_width) * in_height;
    if (npixel % 32 != 0)
        mask.back() = ~((1u << (npixel % 32)) - 1u);
    // On this engine's stream, then synchronised - the default stream is non-blocking, so a NULL-stream
    // copy is not ordered against the kernels that read this mask.
    cuda_err(cudaMemcpyAsync(gpu_res_mask, mask.data(), nwords * sizeof(uint32_t),
                             cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaStreamSynchronize(*stream));
}

void SpotExtractorGPU::SetResolutionMask(const std::vector<uint32_t> &packed_mask) {
    if (packed_mask.size() != nwords)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "SpotExtractorGPU::SetResolutionMask: mask size mismatch");
    cuda_err(cudaMemcpyAsync(gpu_res_mask, packed_mask.data(), nwords * sizeof(uint32_t),
                             cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaStreamSynchronize(*stream));
}

void SpotExtractorGPU::Extract(const uint32_t *gpu_strong, const int32_t *gpu_image,
                               const SpotFindingSettings &settings, std::vector<DiffractionSpot> &spots) {
    const int max_pix = static_cast<int>(settings.max_pix_per_spot);

    count_bits<<<compact_blocks, THREADS, 0, *stream>>>(gpu_strong, gpu_res_mask, gpu_block_count, nwords);
    scan_block_counts<<<1, FINISH_THREADS, 0, *stream>>>(gpu_block_count, gpu_block_offset, gpu_nstrong,
                                                          compact_blocks);
    scatter_bits<<<compact_blocks, 32, 0, *stream>>>(gpu_strong, gpu_res_mask, gpu_block_offset, gpu_image,
                                                      gpu_index, gpu_value, nwords, MAX_STRONG);
    // Fixed grids reading the strong-pixel count from device memory: the host never learns it, so it
    // never has to synchronise in the middle of the frame.
    init_parent<<<512, THREADS, 0, *stream>>>(gpu_parent, gpu_nstrong, MAX_STRONG);
    union_neighbours<<<512, THREADS, 0, *stream>>>(gpu_index, gpu_parent, gpu_nstrong, MAX_STRONG, width);
    resolve_roots<<<512, THREADS, 0, *stream>>>(gpu_parent, gpu_root, gpu_nstrong, MAX_STRONG);
    finish_components<<<1, FINISH_THREADS, 0, *stream>>>(gpu_index, gpu_value, gpu_root, gpu_label,
                                                          gpu_count, gpu_spot, gpu_spot_out, gpu_nspot,
                                                          gpu_nstrong, MAX_STRONG, width, max_pix);
    cuda_err(cudaMemcpyAsync(host_nspot, gpu_nspot, sizeof(uint32_t), cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(host_spot, gpu_spot_out, SPOT_PREFIX * sizeof(SpotExtractorGPUSpot),
                             cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaStreamSynchronize(*stream));      // the only synchronisation in the frame

    const uint32_t nspot = *host_nspot.get();
    const SpotExtractorGPUSpot *s = host_spot.get();
    if (nspot > SPOT_PREFIX) {
        overflow_spot.resize(nspot);
        cuda_err(cudaMemcpyAsync(overflow_spot.data(), gpu_spot_out, nspot * sizeof(SpotExtractorGPUSpot),
                                 cudaMemcpyDeviceToHost, *stream));
        cuda_err(cudaStreamSynchronize(*stream));
        s = overflow_spot.data();
    }

    spots.clear();
    spots.reserve(nspot);
    for (uint32_t i = 0; i < nspot; i++)
        spots.emplace_back(s[i].x, s[i].y, s[i].pixel_count, s[i].photons, s[i].max_photons);
}