#ifndef CORE_RNG_H
#define CORE_RNG_H

#include <cstdint>
#include <random>
#include <vector>
#include <utility>
#include <initializer_list>
#include <type_traits>
#include <limits>

namespace dofs {

class Rng final {
public:
        using engine_type = std::mt19937_64;
        using seed_type   = engine_type::result_type;

        explicit Rng(seed_type seed = default_seed()) noexcept
                : _eng(seed) {}

        void seed(seed_type s) noexcept {
                _eng.seed(s);
        }

        // Real in [0, 1)
        double uniform01() {
                return std::uniform_real_distribution<double> {}(_eng);
        }

        // Integer in [lo, hi) — lo inclusive, hi exclusive
        template <typename Int,
                  typename = std::enable_if_t<std::is_integral<Int>::value>>
        Int uniform_range(Int lo_inclusive, Int hi_exclusive) {
                return std::uniform_int_distribution<Int>(lo_inclusive, hi_exclusive - 1)(_eng);
        }

        // Integer in [0, hi) — convenience overload
        template <typename Int,
                  typename = std::enable_if_t<std::is_integral<Int>::value>>
        Int uniform_range(Int hi_exclusive) {
                return uniform_range<Int>(0, hi_exclusive);
        }

        double uniform_range(double lo_inclusive, double hi_exclusive) {

                return lo_inclusive + (hi_exclusive - lo_inclusive) * uniform01();
        }

        std::uint64_t poisson(double lambda) {
                return static_cast<std::uint64_t>(std::poisson_distribution<std::uint64_t>
                                                  (lambda)(_eng));
        }

        template <typename T>
        T choose_weighted(const std::vector<std::pair<double, T>>& items) {
                return choose_weighted_impl(items.begin(), items.end());
        }

        template <typename T>
        T choose_weighted(std::initializer_list<std::pair<double, T>> items) {
                return choose_weighted_impl(items.begin(), items.end());
        }

private:
        engine_type _eng;

        static constexpr seed_type default_seed() noexcept {
                return 0xC0FFEEULL ^ 0x9E3779B97F4A7C15ULL;
        }

        template <typename Iter>
        auto choose_weighted_impl(Iter first,
                                  Iter last) -> typename std::iterator_traits<Iter>::value_type::second_type {
                using Pair = typename std::iterator_traits<Iter>::value_type;
                using T    = typename Pair::second_type;

                double total = 0.0;

                for (auto it = first; it != last; ++it) {
                        const double w = it->first;

                        if (w > 0.0 && std::isfinite(w))
                                total += w;
                }

                if (!(total > 0.0)) {
                        std::size_t n = 0;

                        for (auto it = first; it != last; ++it)
                                ++n;

                        if (n == 0)
                                return T{};

                        std::size_t idx = static_cast<std::size_t>(uniform_range<std::uint64_t>(0,
                                          static_cast<std::uint64_t>(n)));

                        auto it = first;

                        while (idx--)
                                ++it;

                        return it->second;
                }

                double u = uniform01() * total;

                for (auto it = first; it != last; ++it) {
                        const double w = (it->first > 0.0 &&
                                          std::isfinite(it->first)) ? it->first : 0.0;

                        if (u < w)
                                return it->second;

                        u -= w;
                }

                auto it = last;

                do {
                        --it;
                } while (it->first <= 0.0 || !std::isfinite(it->first));

                return it->second;
        }
};

} // namespace dofs

#endif // CORE_RNG_H
