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

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <mutex>
#include <numeric>
#include <stdexcept>
#include <vector>

#include "../common/ParallelFor.h"

namespace {
    const std::vector<size_t> THREAD_COUNTS = {0, 1, 2, 3, 8, 64};
    const std::vector<int> SIZES = {0, 1, 2, 7, 64, 1000};
}

// The split is fixed and contiguous, which is what lets a pass whose per-element work is independent
// give the serial answer bit for bit. Every element has to land in exactly one slice, the slices have
// to tile [0, n) in order, and there must never be an empty one.
TEST_CASE("ParallelChunks_SlicesTileTheRange", "[ParallelFor]") {
    for (int n: SIZES) {
        for (size_t nthreads: THREAD_COUNTS) {
            CAPTURE(n, nthreads);
            std::mutex m;
            std::vector<std::pair<int, int>> slices;
            std::vector<int> visits(static_cast<size_t>(std::max(n, 1)), 0);

            ParallelChunks(n, nthreads, [&](int lo, int hi) {
                {
                    std::scoped_lock lock(m);
                    slices.emplace_back(lo, hi);
                }
                for (int i = lo; i < hi; i++)
                    visits[static_cast<size_t>(i)]++;   // slices are disjoint, so no lock needed
            });

            std::sort(slices.begin(), slices.end());
            int expected_lo = 0;
            for (const auto &[lo, hi]: slices) {
                CHECK(lo == expected_lo);
                CHECK(hi > lo);                          // never an empty slice
                expected_lo = hi;
            }
            CHECK(expected_lo == n);                     // and they reach the end
            CHECK(slices.size() <= static_cast<size_t>(std::max(n, 0)));
            for (int i = 0; i < n; i++)
                CHECK(visits[static_cast<size_t>(i)] == 1);
        }
    }
}

// Work stealing, so a worker takes whatever is next rather than a fixed slice - but every item still
// has to be done exactly once, whatever order they come out in.
TEST_CASE("ParallelFor_VisitsEveryItemOnce", "[ParallelFor]") {
    for (int n: SIZES) {
        for (size_t nthreads: THREAD_COUNTS) {
            CAPTURE(n, nthreads);
            std::vector<std::atomic<int>> visits(static_cast<size_t>(std::max(n, 1)));
            for (auto &v: visits)
                v.store(0);

            ParallelFor(n, nthreads, [&](int i) { visits[static_cast<size_t>(i)].fetch_add(1); });

            for (int i = 0; i < n; i++)
                CHECK(visits[static_cast<size_t>(i)].load() == 1);
        }
    }
}

// One thread means the caller's loop, in order - the property the serial fallbacks rely on.
TEST_CASE("ParallelFor_IsSerialAndInOrderForOneThread", "[ParallelFor]") {
    for (size_t nthreads: {size_t{0}, size_t{1}}) {
        CAPTURE(nthreads);
        std::vector<int> order;
        ParallelFor(16, nthreads, [&](int i) { order.push_back(i); });
        std::vector<int> expected(16);
        std::iota(expected.begin(), expected.end(), 0);
        CHECK(order == expected);
    }
}

// Splitting the work must not change the answer. Each element is written by exactly one worker, so
// the result has to match the serial loop element for element, at every thread count.
TEST_CASE("ParallelFor_SplitDoesNotChangeTheResult", "[ParallelFor]") {
    constexpr int N = 1000;
    std::vector<double> serial(N);
    for (int i = 0; i < N; i++)
        serial[static_cast<size_t>(i)] = std::sin(i * 0.001) * 1e6 + i;

    for (size_t nthreads: THREAD_COUNTS) {
        CAPTURE(nthreads);
        std::vector<double> chunked(N, 0.0), stolen(N, 0.0);
        ParallelChunks(N, nthreads, [&](int lo, int hi) {
            for (int i = lo; i < hi; i++)
                chunked[static_cast<size_t>(i)] = std::sin(i * 0.001) * 1e6 + i;
        });
        ParallelFor(N, nthreads, [&](int i) {
            stolen[static_cast<size_t>(i)] = std::sin(i * 0.001) * 1e6 + i;
        });
        CHECK(chunked == serial);       // bit for bit, not approximately
        CHECK(stolen == serial);
    }
}

// A negative or zero count is a no-op rather than an error - callers pass a computed size.
TEST_CASE("ParallelFor_DoesNothingForAnEmptyRange", "[ParallelFor]") {
    int calls = 0;
    for (int n: {0, -1, -1000}) {
        ParallelChunks(n, 8, [&](int, int) { calls++; });
        ParallelFor(n, 8, [&](int) { calls++; });
    }
    CHECK(calls == 0);
}

// An exception thrown in a worker reaches the caller rather than terminating: the futures are waited
// on, so the other workers finish first and only then does it propagate.
TEST_CASE("ParallelFor_PropagatesAnException", "[ParallelFor]") {
    CHECK_THROWS_AS(ParallelChunks(64, 4, [](int lo, int) {
        if (lo == 0) throw std::runtime_error("from a chunk");
    }), std::runtime_error);

    CHECK_THROWS_AS(ParallelFor(64, 4, [](int i) {
        if (i == 0) throw std::runtime_error("from an item");
    }), std::runtime_error);
}