// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// random_sampling.hpp
// Standard floating-point (STL std::xxx_distribution) and multiprecision floating-point (sangi::Float):
// a random-number sampling foundation that works with both.
//
// STL's std::uniform_real_distribution<T> etc. require T to be a built-in floating-point type
// (no Float support), so for multiprecision Float we build uniform randoms from 32-bit words
// and implement each distribution via inverse-CDF / Box-Muller / Marsaglia-Tsang, etc.
//
// Each sampleXxxStd branches with if constexpr: standard floating-point uses STL as-is
// (= keeps existing behavior and the random sequence unchanged); otherwise (Float) uses a custom implementation.

#ifndef SANGI_RANDOM_SAMPLING_HPP
#define SANGI_RANDOM_SAMPLING_HPP

#include <random>
#include <cstdint>
#include <type_traits>
#include <math/core/generic_math.hpp>

namespace sangi {
namespace detail {

// Uniform distribution [0, 1). For multiprecision Float, lay out mt19937 32-bit words up to the working precision
// to build the mantissa (entropy >= working precision).
template<typename T>
[[nodiscard]] inline T uniform01(std::mt19937& gen) {
    if constexpr (std::is_floating_point_v<T>) {
        return std::uniform_real_distribution<T>(T(0), T(1))(gen);
    } else {
        const int p = T::defaultPrecision();         // working precision (bits)
        const int words = p / 32 + 2;                 // number of 32-bit words (with margin)
        const T inv = T(1) / T(static_cast<std::uint64_t>(1) << 32);  // 2^-32 (exact)
        T result(0);
        T scale(1);
        for (int i = 0; i < words; ++i) {
            scale = scale * inv;
            result = result + T(static_cast<std::uint64_t>(gen())) * scale;
        }
        return result;
    }
}

// Uniform distribution (0, 1] (excludes 0 so it can be used as a log argument).
template<typename T>
[[nodiscard]] inline T uniform01_pos(std::mt19937& gen) {
    return T(1) - uniform01<T>(gen);
}

// Standard normal distribution N(0,1). Box-Muller method (Float).
template<typename T>
[[nodiscard]] inline T sampleNormalStd(std::mt19937& gen) {
    if constexpr (std::is_floating_point_v<T>) {
        return std::normal_distribution<T>(T(0), T(1))(gen);
    } else {
        T u1 = uniform01_pos<T>(gen);                 // (0,1]
        T u2 = uniform01<T>(gen);
        T r = generic_sqrt(T(-2) * generic_log(u1));
        T theta = T(2) * generic_pi<T>() * u2;
        return r * generic_cos(theta);
    }
}

// Standard exponential distribution (rate 1). Inverse CDF: -ln(1-u).
template<typename T>
[[nodiscard]] inline T sampleExponentialStd(std::mt19937& gen) {
    if constexpr (std::is_floating_point_v<T>) {
        return std::exponential_distribution<T>(T(1))(gen);
    } else {
        return -generic_log(uniform01_pos<T>(gen));
    }
}

// Standard gamma distribution (shape a, scale 1). Marsaglia-Tsang method (Float).
template<typename T>
[[nodiscard]] inline T sampleGammaStd(std::mt19937& gen, T a) {
    if constexpr (std::is_floating_point_v<T>) {
        return std::gamma_distribution<T>(a, T(1))(gen);
    } else {
        if (a < T(1)) {
            // a<1 is corrected by multiplying the Gamma(a+1) sample by U^{1/a} (Marsaglia-Tsang boost)
            T u = uniform01_pos<T>(gen);
            return sampleGammaStd<T>(gen, a + T(1)) * generic_pow(u, T(1) / a);
        }
        const T d = a - T(1) / T(3);
        const T c = T(1) / generic_sqrt(T(9) * d);
        for (;;) {
            T x = sampleNormalStd<T>(gen);
            T v = T(1) + c * x;
            if (v <= T(0)) continue;
            v = v * v * v;
            T u = uniform01<T>(gen);
            T x2 = x * x;
            if (u < T(1) - T(0.0331) * x2 * x2) return d * v;
            if (generic_log(u) < T(1) / T(2) * x2 + d * (T(1) - v + generic_log(v)))
                return d * v;
        }
    }
}

// Sample the Poisson distribution (mean lambda) by Knuth's method (for Float; returns a count as int).
// For large lambda this takes O(lambda) iterations but is correct.
template<typename T>
[[nodiscard]] inline int samplePoissonKnuth(std::mt19937& gen, T lambda) {
    const T L = generic_exp(-lambda);
    int k = 0;
    T p(1);
    do {
        ++k;
        p = p * uniform01<T>(gen);
    } while (p > L);
    return k - 1;
}

} // namespace detail
} // namespace sangi

#endif // SANGI_RANDOM_SAMPLING_HPP