// 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.

// 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();
}