// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later
//
// normal.hpp
//
// Fast sampling of the normal distribution:
// - Xoshiro256pp                 — fast PRNG (xoshiro256++ by Blackman & Vigna 2018)
// - NormalGenerator<T>           — Ziggurat method (Marsaglia & Tsang 2000, N=256)
// - MarsagliaPolarGenerator<T>   — Marsaglia polar method (no sin/cos)
// - BoxMullerGenerator<T>        — Box-Muller method (for teaching/verification)
//
// All provide next() and fill(Container). next() returns one sample;
// fill() writes a normal-distribution sample into each element of Container.
//
// Reference: web/note/signal-processing/signal-generation/noise/white-noise.html

#ifndef SANGI_RANDOM_NORMAL_HPP
#define SANGI_RANDOM_NORMAL_HPP

#include <array>
#include <bit>      // std::rotl
#include <cmath>
#include <cstdint>
#include <numbers>
#include <type_traits>

namespace sangi {
namespace random {

// =====================================================================
// Xoshiro256pp — xoshiro256++ PRNG
//   Blackman & Vigna (2018), https://prng.di.unimi.it/
// 256-bit state, period 2^256 - 1, faster than std::mt19937 with good statistical quality.
// =====================================================================

class Xoshiro256pp {
public:
    using result_type = std::uint64_t;
    static constexpr result_type min() noexcept { return 0; }
    static constexpr result_type max() noexcept { return ~result_type(0); }

    explicit Xoshiro256pp(std::uint64_t seed = 0x9E3779B97F4A7C15ULL) noexcept {
        seed_state(seed ? seed : 0x9E3779B97F4A7C15ULL);
    }

    void seed(std::uint64_t s) noexcept { seed_state(s ? s : 0x9E3779B97F4A7C15ULL); }

    result_type operator()() noexcept {
        const std::uint64_t result = std::rotl(state_[0] + state_[3], 23) + state_[0];
        const std::uint64_t t = state_[1] << 17;
        state_[2] ^= state_[0];
        state_[3] ^= state_[1];
        state_[1] ^= state_[2];
        state_[0] ^= state_[3];
        state_[2] ^= t;
        state_[3] = std::rotl(state_[3], 45);
        return result;
    }

private:
    std::uint64_t state_[4];

    // Expand the seed into a 256-bit state with SplitMix64
    void seed_state(std::uint64_t s) noexcept {
        for (int i = 0; i < 4; ++i) {
            s += 0x9E3779B97F4A7C15ULL;
            std::uint64_t z = s;
            z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
            z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
            z = z ^ (z >> 31);
            state_[i] = z;
        }
    }
};

// =====================================================================
// detail — Ziggurat table (N = 256, built in double precision)
// =====================================================================

namespace detail {

inline constexpr std::size_t ZIG_N = 256;

// Right half of the standard-normal PDF: phi(x) = exp(-x^2/2)  (normalization constant may be omitted)
inline double zig_phi(double x) noexcept {
    return std::exp(-0.5 * x * x);
}

struct ZigguratTable {
    // x[i]: right-edge x coordinate of each layer (x[0] = R = the tail entry, as a remnant of x_{N-1};
    //       in this implementation the array layout has x[0]=R and x[1]..x[N-1] as layer right edges)
    std::array<double, ZIG_N + 1> x{};
    // y[i] = phi(x[i])
    std::array<double, ZIG_N + 1> y{};
    // ratio: x[i-1] / x[i] for fast acceptance. Kept as double rather than converted to integer comparison
    //        (prioritizing testability / portability)
    std::array<double, ZIG_N> ratio{};

    // R = x_{N-1}: tail entry (Doornik 2005 value ~ 3.6541528853610088)
    double R{};
    // V: area of each layer (value of R * phi(R) + tail_area)
    double V{};
};

// Table construction:
//   equation: V = R * phi(R) + ∫_R^∞ phi(t) dt
//   R is solved by Newton's method / bisection. Here the standard R value for N=256
//   is fixed in advance (Doornik 2005); V is computed and x[i] generated by the recurrence.
inline const ZigguratTable& ziggurat_table() {
    static const ZigguratTable table = []() {
        ZigguratTable t{};
        // Doornik (2005) "An improved Ziggurat method..." N=256
        // R, V are literature values (note: phi normalization omitted, phi(x) = exp(-x^2/2))
        // ∫_R^∞ exp(-t^2/2) dt = sqrt(pi/2) * erfc(R/sqrt(2))
        t.R = 3.6541528853610088;
        const double pi = std::numbers::pi_v<double>;
        const double tail = std::sqrt(pi / 2.0)
                            * std::erfc(t.R / std::sqrt(2.0));
        t.V = t.R * zig_phi(t.R) + tail;

        // x[N] = 0 (the right edge of the bottom layer is not 0; rather, layer indices 1..N-1 are
        // rectangles and layer 0 is the tail). Here the array layout is as follows:
        //   x[0]   = R                 (tail entry, fast-acceptance threshold)
        //   x[i]   = right edge of the i-th layer  (i = 1..N-1)
        //   x[N]   = 0                 (sentinel, end of the recurrence)
        //   y[i]   = phi(x[i])
        // recurrence (backward from i = N-1): x[i-1] = phi^{-1}( y[i] + V / x[i] )
        //   where y[i] = phi(x[i]); add V/x[i] to form y[i-1].
        //   phi^{-1}(u) = sqrt(-2 ln(u))
        t.x[0] = t.R;
        t.y[0] = zig_phi(t.R);
        // x at i = 1: one convention writes the top layer as x[1] = x_{N-1} = R (= outermost rectangle), and
        // another writes x[1] as the bottom layer (the small rectangle closest to R).
        // Here we use "the larger i, the smaller x[i], with i=N-1 near the origin":
        //   x[1] = R (top layer = right edge of the outermost rectangle = layer just below the tail entry)
        //   x[i+1] = phi^{-1}( phi(x[i]) + V / x[i] )
        // y[i] = phi(x[i])
        t.x[1] = t.R;
        t.y[1] = zig_phi(t.R);
        for (std::size_t i = 1; i < ZIG_N - 1; ++i) {
            double u = t.y[i] + t.V / t.x[i];
            // u becomes phi(x[i+1]). u > 1 is numerically impossible, but clamp anyway.
            if (u > 1.0) u = 1.0;
            t.x[i + 1] = std::sqrt(-2.0 * std::log(u));
            t.y[i + 1] = u;
        }
        t.x[ZIG_N] = 0.0;
        t.y[ZIG_N] = 1.0;  // phi(0) = 1 (normalization constant omitted)

        // ratio[i] = x[i+1] / x[i] (i = 0..N-1)
        // used for the test |U1| < ratio[i], equivalent to the fast-acceptance |U1*x[i]| < x[i+1]
        for (std::size_t i = 0; i < ZIG_N; ++i) {
            t.ratio[i] = (t.x[i] > 0.0) ? (t.x[i + 1] / t.x[i]) : 0.0;
        }
        return t;
    }();
    return table;
}

// Extract one 64-bit integer from any URBG
template <typename URBG>
inline std::uint64_t next_uint64(URBG& gen) noexcept {
    if constexpr (sizeof(typename URBG::result_type) >= 8) {
        return static_cast<std::uint64_t>(gen());
    } else {
        const std::uint64_t hi = static_cast<std::uint64_t>(gen());
        const std::uint64_t lo = static_cast<std::uint64_t>(gen());
        return (hi << 32) | (lo & 0xFFFFFFFFULL);
    }
}

// Uniform double random in [0, 1) (built from the top 53 bits)
template <typename URBG>
inline double uniform01(URBG& gen) noexcept {
    std::uint64_t u = next_uint64(gen) >> 11;  // 53 bit
    return static_cast<double>(u) * (1.0 / static_cast<double>(1ULL << 53));
}

// Generate standard normal N(0, 1) by the Ziggurat method
template <typename URBG>
inline double ziggurat_normal(URBG& gen) noexcept {
    const ZigguratTable& T = ziggurat_table();
    for (int iter = 0; iter < 1000; ++iter) {  // normally finishes in 1-2 iterations
        const std::uint64_t r = next_uint64(gen);
        // top 8 bits = layer index (0..255), bottom 1 bit = sign, the rest = U1
        const std::size_t i = (r >> 56) & 0xFFu;
        const int sign = (static_cast<int>(r) & 1) ? 1 : -1;
        // U1 is converted to a 53-bit double (from the remaining 55 bits after consuming the sign bit)
        const std::uint64_t u_bits = (r >> 1) & ((1ULL << 53) - 1);
        const double U1 = static_cast<double>(u_bits) / static_cast<double>(1ULL << 53);

        // candidate x = U1 * x[i]
        const double x_cand = U1 * T.x[i];

        // Step 3: fast acceptance (U1 < ratio[i])
        if (U1 < T.ratio[i]) {
            return sign * x_cand;
        }

        if (i == 0) {
            // Step 5: tail sampling (Marsaglia's exponential method)
            double xt, yt;
            do {
                const double u_a = uniform01(gen);
                const double u_b = uniform01(gen);
                // guard because log becomes -inf when u_a, u_b are extremely close to 0
                const double la = (u_a > 0.0) ? std::log(u_a) : -700.0;
                const double lb = (u_b > 0.0) ? std::log(u_b) : -700.0;
                xt = -la / T.R;
                yt = -lb;
            } while (2.0 * yt <= xt * xt);
            return sign * (T.R + xt);
        }

        // Step 4: wedge comparison
        const double y_cand = T.y[i] + uniform01(gen) * (T.y[i - 1] - T.y[i]);
        if (y_cand < zig_phi(x_cand)) {
            return sign * x_cand;
        }
        // else: reject, continue the loop
    }
    // not exiting within 1000 iterations is numerically impossible (upper-bound guard)
    return 0.0;
}

}  // namespace detail

// =====================================================================
// NormalGenerator<T> — N(μ, σ) by the Ziggurat method
// =====================================================================

/**
 * @brief Fast normal-distribution generator (Ziggurat method, N=256)
 *
 * Over 95% of samples need only "table lookup + integer comparison",
 * maximizing throughput for bulk generation. The internal PRNG is xoshiro256++.
 *
 * @tparam T sample type (float / double expected; computed in double precision then cast)
 *
 * Example usage:
 *   sangi::random::NormalGenerator<float> g(0.0f, 1.0f, 42);
 *   float x = g.next();
 *   std::vector<float> buf(44100);
 *   g.fill(buf);
 */
template<typename T = double>
class NormalGenerator {
    static_assert(std::is_floating_point_v<T>,
                  "NormalGenerator: T must be a floating point type");
public:
    using value_type = T;

    explicit NormalGenerator(T mu = T{0}, T sigma = T{1},
                             std::uint64_t seed = 0x9E3779B97F4A7C15ULL)
        : mu_(mu), sigma_(sigma), gen_(seed) {}

    /// Generate one sample
    T next() noexcept {
        return static_cast<T>(mu_ + sigma_ * static_cast<T>(detail::ziggurat_normal(gen_)));
    }

    /// Write samples into the entire container (scalar loop internally; SIMD planned)
    template<typename Container>
    void fill(Container& c) noexcept {
        for (auto& v : c) v = next();
    }

    /// Reseed
    void seed(std::uint64_t s) noexcept { gen_.seed(s); }

    /// Parameter access
    T mu() const noexcept { return mu_; }
    T sigma() const noexcept { return sigma_; }

private:
    T mu_;
    T sigma_;
    Xoshiro256pp gen_;
};

// =====================================================================
// MarsagliaPolarGenerator<T> — Marsaglia polar method
//   no sin/cos; generates N(0,1) two at a time using sqrt + log.
//   acceptance rate π/4 ≈ 78.5%; one is cached and reused.
// =====================================================================

template<typename T = double>
class MarsagliaPolarGenerator {
    static_assert(std::is_floating_point_v<T>,
                  "MarsagliaPolarGenerator: T must be a floating point type");
public:
    using value_type = T;

    explicit MarsagliaPolarGenerator(T mu = T{0}, T sigma = T{1},
                                     std::uint64_t seed = 0x9E3779B97F4A7C15ULL)
        : mu_(mu), sigma_(sigma), gen_(seed) {}

    T next() noexcept {
        if (has_cached_) {
            has_cached_ = false;
            return static_cast<T>(mu_ + sigma_ * cached_);
        }
        T u, v, s;
        do {
            u = T(2) * static_cast<T>(detail::uniform01(gen_)) - T(1);
            v = T(2) * static_cast<T>(detail::uniform01(gen_)) - T(1);
            s = u * u + v * v;
        } while (s >= T(1) || s == T(0));
        const T factor = std::sqrt(T(-2) * std::log(s) / s);
        cached_ = v * factor;
        has_cached_ = true;
        return static_cast<T>(mu_ + sigma_ * u * factor);
    }

    template<typename Container>
    void fill(Container& c) noexcept {
        for (auto& x : c) x = next();
    }

    void seed(std::uint64_t s) noexcept {
        gen_.seed(s);
        has_cached_ = false;
    }

    T mu() const noexcept { return mu_; }
    T sigma() const noexcept { return sigma_; }

private:
    T mu_;
    T sigma_;
    Xoshiro256pp gen_;
    T cached_{};
    bool has_cached_ = false;
};

// =====================================================================
// BoxMullerGenerator<T> — Box-Muller method (for teaching/verification)
//   acceptance rate 100%; calls sin/cos once per pair.
// =====================================================================

template<typename T = double>
class BoxMullerGenerator {
    static_assert(std::is_floating_point_v<T>,
                  "BoxMullerGenerator: T must be a floating point type");
public:
    using value_type = T;

    explicit BoxMullerGenerator(T mu = T{0}, T sigma = T{1},
                                std::uint64_t seed = 0x9E3779B97F4A7C15ULL)
        : mu_(mu), sigma_(sigma), gen_(seed) {}

    T next() noexcept {
        if (has_cached_) {
            has_cached_ = false;
            return static_cast<T>(mu_ + sigma_ * cached_);
        }
        // constrain to >= ε because log becomes -inf when U1 is 0
        T u1, u2;
        do {
            u1 = static_cast<T>(detail::uniform01(gen_));
        } while (u1 <= T(0));
        u2 = static_cast<T>(detail::uniform01(gen_));
        const T r = std::sqrt(T(-2) * std::log(u1));
        const T theta = T(2) * std::numbers::pi_v<T> * u2;
        cached_ = r * std::sin(theta);
        has_cached_ = true;
        return static_cast<T>(mu_ + sigma_ * r * std::cos(theta));
    }

    template<typename Container>
    void fill(Container& c) noexcept {
        for (auto& x : c) x = next();
    }

    void seed(std::uint64_t s) noexcept {
        gen_.seed(s);
        has_cached_ = false;
    }

    T mu() const noexcept { return mu_; }
    T sigma() const noexcept { return sigma_; }

private:
    T mu_;
    T sigma_;
    Xoshiro256pp gen_;
    T cached_{};
    bool has_cached_ = false;
};

}  // namespace random
}  // namespace sangi

#endif  // SANGI_RANDOM_NORMAL_HPP