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

#pragma once

#include <algorithm>
#include <atomic>
#include <future>
#include <vector>

// Two shapes of "run this over a range on several threads", used by the analysis code. Both take the
// worker count from the caller rather than asking the hardware, so a run that was told how many
// threads to use keeps to it.

// How many workers a pass over `n` cheap items should use: enough that each gets at least
// `min_per_thread` of them, and never more than the caller was given. Both helpers below start a
// thread per chunk, which costs tens of microseconds, so a pass that is repeated - or one whose data
// is small - can easily spend more on splitting the work than on doing it. That is not a large
// machine's problem: it is what makes the same code right on an 8-core laptop, a 16-core desktop and
// a two-socket node, none of which should be handed 48 chunks of a few thousand items.
inline size_t ThreadsForWork(size_t n, size_t nthreads, size_t min_per_thread = 32768) {
    if (nthreads <= 1 || n == 0) return 1;
    return std::clamp<size_t>(n / min_per_thread, 1, nthreads);
}

// Chunked: each worker gets one contiguous [lo, hi) range and there is no per-item synchronisation.
// Right for millions of cheap uniform items - the CPU stand-in for a flat CUDA grid-stride kernel.
// The split is fixed and deterministic, so a pass whose per-element work is independent gives the
// same answer as the serial loop, bit for bit.
template <typename Fn>
void ParallelChunks(int n, size_t nthreads, Fn fn) {
    if (n <= 0) return;
    const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
    if (nt == 1) { fn(0, n); return; }
    const int chunk = (n + nt - 1) / nt;
    std::vector<std::future<void>> futures;
    futures.reserve(nt);
    for (int t = 0; t < nt; t++) {
        const int lo = t * chunk, hi = std::min(n, lo + chunk);
        if (lo >= hi) break;
        futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); }));
    }
    for (auto &f : futures) f.get();
}

// Work-stealing per item, off a shared atomic counter: one atomic per item, so use it only where the
// per-item work is heavy and uneven (per-frame fits, per-ring selections) and the atomic amortises.
// For millions of tiny uniform items a per-item atomic is pure contention - use ParallelChunks.
template <typename Fn>
void ParallelFor(int n, size_t nthreads, Fn fn) {
    if (n <= 0) return;
    if (nthreads <= 1 || n == 1) {
        for (int i = 0; i < n; i++) fn(i);
        return;
    }
    const size_t local = std::min(nthreads, static_cast<size_t>(n));
    std::atomic<int> next = 0;
    std::vector<std::future<void>> futures;
    futures.reserve(local);
    for (size_t t = 0; t < local; t++)
        futures.emplace_back(std::async(std::launch::async, [&] {
            for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
                fn(i);
        }));
    for (auto &f : futures) f.get();
}