#pragma once

#include <algorithm>
#include <chrono>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <mutex>
#include <string>
#include <vector>
#include <mpi.h>

namespace mean_field::profiling {
    struct Statistics {
        unsigned long long observations{0};
        unsigned long long warmups{0};
        unsigned long long samples{0};
        unsigned long long warmup_target{0};
        double total_seconds{0.0};
        double minimum_seconds{std::numeric_limits<double>::infinity()};
        double maximum_seconds{0.0};
    };

    class Registry {
    public:
        static Registry& Get() {
            static Registry registry;
            return registry;
        }

        void Record(const std::string& label, const double seconds, const unsigned long long warmup_count) {
            std::scoped_lock lock(m_mutex);
            Statistics& statistics = m_statistics[label];

            statistics.warmup_target = std::max(statistics.warmup_target, warmup_count);
            const bool is_warmup = statistics.observations < statistics.warmup_target;
            ++statistics.observations;

            if (is_warmup) {
                ++statistics.warmups;
                return;
            }

            ++statistics.samples;
            statistics.total_seconds += seconds;
            statistics.minimum_seconds = std::min(statistics.minimum_seconds, seconds);
            statistics.maximum_seconds = std::max(statistics.maximum_seconds, seconds);
        }

        void Reset() {
            std::scoped_lock lock(m_mutex);
            m_statistics.clear();
        }

        void Print(MPI_Comm communicator) const {
            const std::map<std::string, Statistics> snapshot = GetSnapshot();

            int mpi_initialized = 0;
            int mpi_finalized = 0;
            MPI_Initialized(&mpi_initialized);
            if (mpi_initialized) MPI_Finalized(&mpi_finalized);

            const bool use_mpi = mpi_initialized && !mpi_finalized;
            int rank = 0;
            int communicator_size = 1;

            if (use_mpi) {
                MPI_Comm_rank(communicator, &rank);
                MPI_Comm_size(communicator, &communicator_size);
            }

            if (rank == 0) {
                std::cout << '\n';
                std::cout << std::left << std::setw(42) << "Profile Region"
                          << std::right << std::setw(11) << "Samples"
                          << std::setw(10) << "Warmups"
                          << std::setw(14) << "Avg Max ms"
                          << std::setw(14) << "Min ms"
                          << std::setw(14) << "Max ms"
                          << std::setw(14) << "Total Max s" << '\n';
                std::cout << std::string(119, '-') << '\n';
            }

            for (const auto& [label, local_statistics] : snapshot) {
                unsigned long long minimum_samples = local_statistics.samples;
                unsigned long long maximum_samples = local_statistics.samples;
                unsigned long long maximum_warmups = local_statistics.warmups;

                double local_average = local_statistics.samples > 0 ? local_statistics.total_seconds / static_cast<double>(local_statistics.samples) : 0.0;
                double local_minimum = local_statistics.samples > 0 ? local_statistics.minimum_seconds : std::numeric_limits<double>::infinity();
                double local_maximum = local_statistics.maximum_seconds;
                double local_total = local_statistics.total_seconds;

                double maximum_rank_average = local_average;
                double global_minimum = local_minimum;
                double global_maximum = local_maximum;
                double maximum_rank_total = local_total;

                if (use_mpi) {
                    MPI_Allreduce(&local_statistics.samples, &minimum_samples, 1, MPI_UNSIGNED_LONG_LONG, MPI_MIN, communicator);
                    MPI_Allreduce(&local_statistics.samples, &maximum_samples, 1, MPI_UNSIGNED_LONG_LONG, MPI_MAX, communicator);
                    MPI_Allreduce(&local_statistics.warmups, &maximum_warmups, 1, MPI_UNSIGNED_LONG_LONG, MPI_MAX, communicator);
                    MPI_Allreduce(&local_average, &maximum_rank_average, 1, MPI_DOUBLE, MPI_MAX, communicator);
                    MPI_Allreduce(&local_minimum, &global_minimum, 1, MPI_DOUBLE, MPI_MIN, communicator);
                    MPI_Allreduce(&local_maximum, &global_maximum, 1, MPI_DOUBLE, MPI_MAX, communicator);
                    MPI_Allreduce(&local_total, &maximum_rank_total, 1, MPI_DOUBLE, MPI_MAX, communicator);
                }

                if (!std::isfinite(global_minimum)) global_minimum = 0.0;

                if (rank == 0) {
                    const std::string sample_string = minimum_samples == maximum_samples
                        ? std::to_string(minimum_samples)
                        : std::to_string(minimum_samples) + "-" + std::to_string(maximum_samples);

                    std::cout << std::left << std::setw(100) << label
                              << std::right << std::setw(11) << sample_string
                              << std::setw(10) << maximum_warmups
                              << std::setw(14) << std::fixed << std::setprecision(3) << 1.0e3 * maximum_rank_average
                              << std::setw(14) << 1.0e3 * global_minimum
                              << std::setw(14) << 1.0e3 * global_maximum
                              << std::setw(14) << std::setprecision(6) << maximum_rank_total << '\n';
                }
            }

            if (rank == 0) {
                std::cout << std::string(119, '=') << '\n';
                std::cout << "MPI ranks: " << communicator_size << "\n\n";
            }
        }

    private:
        [[nodiscard]] std::map<std::string, Statistics> GetSnapshot() const {
            std::scoped_lock lock(m_mutex);
            return m_statistics;
        }

    private:
        mutable std::mutex m_mutex;
        std::map<std::string, Statistics> m_statistics;
    };

    class ScopedTimer {
    public:
        ScopedTimer(std::string label, const unsigned long long warmup_count)
            : m_label(std::move(label)),
              m_warmup_count(warmup_count),
              m_start(std::chrono::steady_clock::now()) {}

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

        ~ScopedTimer() {
            try {
                const auto stop = std::chrono::steady_clock::now();
                const double seconds = std::chrono::duration<double>(stop - m_start).count();
                Registry::Get().Record(m_label, seconds, m_warmup_count);
            } catch (...) {}
        }

    private:
        std::string m_label;
        unsigned long long m_warmup_count;
        std::chrono::steady_clock::time_point m_start;
    };
}

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

#define MEAN_FIELD_PROFILE_SCOPE_WARMUP(label, warmup_count) \
    ::mean_field::profiling::ScopedTimer MEAN_FIELD_PROFILE_JOIN(mean_field_profile_timer_, __COUNTER__)(label, warmup_count)

#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_RESET() \
    ::mean_field::profiling::Registry::Get().Reset()

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