#pragma once

#include <chrono>
#include <cstddef>
#include <cstdint>
#include <iosfwd>
#include <map>
#include <memory>
#include <string>
#include <string_view>
#include <vector>

#include <mpi.h>

#ifndef MEAN_FIELD_ENABLE_PROFILING
#define MEAN_FIELD_ENABLE_PROFILING 0
#endif

namespace mean_field::profiling {
    struct Statistics {
        std::uint64_t observations{0};
        std::uint64_t warmups{0};
        std::uint64_t samples{0};
        std::uint64_t warmup_target{0};
        std::uint64_t work_units{0};
        double total_seconds{0.0};
        double minimum_seconds{0.0};
        double maximum_seconds{0.0};
    };

    struct DistributedStatistics {
        std::string label;
        std::uint64_t minimum_samples{0};
        std::uint64_t maximum_samples{0};
        std::uint64_t maximum_warmups{0};
        std::uint64_t minimum_work_units{0};
        std::uint64_t maximum_work_units{0};
        double maximum_rank_average_seconds{0.0};
        double global_minimum_seconds{0.0};
        double global_maximum_seconds{0.0};
        double maximum_rank_total_seconds{0.0};
    };

    class Registry {
    public:
        static Registry &Get();

        Registry(const Registry &)            = delete;
        Registry &operator=(const Registry &) = delete;
        Registry(Registry &&)                 = delete;
        Registry &operator=(Registry &&)      = delete;

        ~Registry();

        void Record(
            std::string_view label,
            double seconds,
            std::uint64_t warmup_count = 0
        );

        void AddCount(
            std::string_view label,
            std::uint64_t work_units
        );

        void Reset();

        [[nodiscard]] std::map<
            std::string,
            Statistics,
            std::less<>>
        Snapshot() const;

        [[nodiscard]] std::vector<DistributedStatistics> Aggregate(MPI_Comm communicator) const;

        void Print(
            MPI_Comm communicator,
            std::ostream &stream
        ) const;

        void Print(MPI_Comm communicator) const;

        void PrintCsv(
            MPI_Comm communicator,
            std::ostream &stream
        ) const;

    private:
        friend class Region;

        Registry();

        [[nodiscard]] std::size_t Register(
            std::string_view label,
            std::uint64_t warmup_count
        );

        void Record(
            std::size_t region,
            double seconds
        ) noexcept;

        void AddCount(
            std::size_t region,
            std::uint64_t work_units
        ) noexcept;

        struct Impl;
        std::unique_ptr<Impl> m_impl;
    };

    class Region {
    public:
        explicit Region(
            std::string_view label,
            std::uint64_t warmup_count = 0
        );

        void Record(double seconds) const noexcept;
        void AddCount(std::uint64_t work_units) const noexcept;

    private:
        std::size_t m_region;
    };

    class ScopedTimer {
    public:
        explicit ScopedTimer(const Region &region) noexcept;

        ScopedTimer(const ScopedTimer &)            = delete;
        ScopedTimer &operator=(const ScopedTimer &) = delete;
        ScopedTimer(ScopedTimer &&)                 = delete;
        ScopedTimer &operator=(ScopedTimer &&)      = delete;

        ~ScopedTimer() noexcept;

    private:
        const Region &m_region;
        std::chrono::steady_clock::time_point m_start;
    };
} // namespace mean_field::profiling

#define MEAN_FIELD_PROFILE_JOIN_IMPL(left, right) left##right
#define MEAN_FIELD_PROFILE_JOIN(left, right) MEAN_FIELD_PROFILE_JOIN_IMPL(left, right)

#if MEAN_FIELD_ENABLE_PROFILING

#define MEAN_FIELD_PROFILE_SCOPE_IMPL(label, warmup_count, identifier)                                                 \
    static const ::mean_field::profiling::Region MEAN_FIELD_PROFILE_JOIN(mean_field_profile_region_, identifier)(      \
        label, warmup_count                                                                                            \
    );                                                                                                                 \
    const ::mean_field::profiling::ScopedTimer MEAN_FIELD_PROFILE_JOIN(mean_field_profile_timer_, identifier)(         \
        MEAN_FIELD_PROFILE_JOIN(mean_field_profile_region_, identifier)                                                \
    )

#define MEAN_FIELD_PROFILE_SCOPE_WARMUP(label, warmup_count)                                                           \
    MEAN_FIELD_PROFILE_SCOPE_IMPL(label, warmup_count, __COUNTER__)

#define MEAN_FIELD_PROFILE_SCOPE(label) MEAN_FIELD_PROFILE_SCOPE_WARMUP(label, 1)

#define MEAN_FIELD_PROFILE_CALL_WARMUP(label, warmup_count, ...)                                                       \
    do {                                                                                                               \
        MEAN_FIELD_PROFILE_SCOPE_WARMUP(label, warmup_count);                                                          \
        __VA_ARGS__;                                                                                                   \
    } while (false)

#define MEAN_FIELD_PROFILE_CALL(label, ...) MEAN_FIELD_PROFILE_CALL_WARMUP(label, 1, __VA_ARGS__)

#define MEAN_FIELD_PROFILE_EVALUATE_IMPL(label, warmup_count, identifier, ...)                                         \
    ([&]() -> decltype(auto) {                                                                                         \
        MEAN_FIELD_PROFILE_SCOPE_IMPL(label, warmup_count, identifier);                                                \
        return (__VA_ARGS__);                                                                                          \
    }())

#define MEAN_FIELD_PROFILE_EVALUATE_WARMUP(label, warmup_count, ...)                                                   \
    MEAN_FIELD_PROFILE_EVALUATE_IMPL(label, warmup_count, __COUNTER__, __VA_ARGS__)

#define MEAN_FIELD_PROFILE_EVALUATE(label, ...) MEAN_FIELD_PROFILE_EVALUATE_WARMUP(label, 1, __VA_ARGS__)

#define MEAN_FIELD_PROFILE_COUNT_IMPL(label, work_units, identifier)                                                   \
    do {                                                                                                               \
        static const ::mean_field::profiling::Region MEAN_FIELD_PROFILE_JOIN(mean_field_profile_counter_, identifier)( \
            label                                                                                                      \
        );                                                                                                             \
        MEAN_FIELD_PROFILE_JOIN(mean_field_profile_counter_, identifier).AddCount(work_units);                         \
    } while (false)

#define MEAN_FIELD_PROFILE_COUNT(label, work_units) MEAN_FIELD_PROFILE_COUNT_IMPL(label, work_units, __COUNTER__)

#define MEAN_FIELD_PROFILE_RESET() ::mean_field::profiling::Registry::Get().Reset()

#define MEAN_FIELD_PROFILE_PRINT(communicator) ::mean_field::profiling::Registry::Get().Print(communicator)

#define MEAN_FIELD_PROFILE_PRINT_CSV(communicator, stream)                                                             \
    ::mean_field::profiling::Registry::Get().PrintCsv(communicator, stream)

#else

#define MEAN_FIELD_PROFILE_SCOPE_WARMUP(label, warmup_count) ((void)0)
#define MEAN_FIELD_PROFILE_SCOPE(label) ((void)0)

#define MEAN_FIELD_PROFILE_CALL_WARMUP(label, warmup_count, ...)                                                       \
    do {                                                                                                               \
        __VA_ARGS__;                                                                                                   \
    } while (false)

#define MEAN_FIELD_PROFILE_CALL(label, ...) MEAN_FIELD_PROFILE_CALL_WARMUP(label, 1, __VA_ARGS__)

#define MEAN_FIELD_PROFILE_EVALUATE_WARMUP(label, warmup_count, ...) (__VA_ARGS__)
#define MEAN_FIELD_PROFILE_EVALUATE(label, ...) (__VA_ARGS__)
#define MEAN_FIELD_PROFILE_COUNT(label, work_units) ((void)0)
#define MEAN_FIELD_PROFILE_RESET() ((void)0)
#define MEAN_FIELD_PROFILE_PRINT(communicator) ((void)0)
#define MEAN_FIELD_PROFILE_PRINT_CSV(communicator, stream) ((void)0)

#endif