#include "profile.h"

#include <algorithm>
#include <catch2/catch_test_macros.hpp>
#include <sstream>
#include <string>
#include <vector>

namespace {
    const mean_field::profiling::DistributedStatistics *find_region(
        const std::vector<mean_field::profiling::DistributedStatistics> &statistics,
        const std::string &label
    ) {
        const auto iterator =
            std::ranges::find(statistics, label, &mean_field::profiling::DistributedStatistics::label);
        return iterator != statistics.end() ? &*iterator : nullptr;
    }
} // namespace

TEST_CASE(
    "MPI Profiling Aggregates Rank-Local Label Sets Without Collective Divergence",
    "[mpi][profiling][distributed]"
) {
    int rank = 0;
    int size = 1;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    mean_field::profiling::Registry &registry = mean_field::profiling::Registry::Get();
    registry.Reset();
    registry.Record("common-region", static_cast<double>(rank + 1));
    registry.AddCount("common-region", static_cast<std::uint64_t>(10 + rank));

    const std::string local_label = "rank-" + std::to_string(rank) + "-only";
    registry.Record(local_label, 0.125 * static_cast<double>(rank + 1));

    const auto aggregate = registry.Aggregate(MPI_COMM_WORLD);
    const auto *common   = find_region(aggregate, "common-region");
    CHECK(common != nullptr);
    if (common != nullptr) {
        CHECK(common->minimum_samples == 1);
        CHECK(common->maximum_samples == 1);
        CHECK(common->minimum_work_units == 10);
        CHECK(common->maximum_work_units == static_cast<std::uint64_t>(9 + size));
        CHECK(common->global_minimum_seconds == 1.0);
        CHECK(common->global_maximum_seconds == static_cast<double>(size));
    }

    for (int owner = 0; owner < size; ++owner) {
        const auto *local = find_region(aggregate, "rank-" + std::to_string(owner) + "-only");
        CHECK(local != nullptr);
        if (local != nullptr) {
            CHECK(local->minimum_samples == 0);
            CHECK(local->maximum_samples == 1);
        }
    }

    std::ostringstream csv;
    registry.PrintCsv(MPI_COMM_WORLD, csv);
    if (rank == 0) {
        CHECK(csv.str().find("common-region") != std::string::npos);
        CHECK(csv.str().find("," + std::to_string(size) + "\n") != std::string::npos);
    } else {
        CHECK(csv.str().empty());
    }
}