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

// quasi_random.hpp
// Quasi-random sequences (low-discrepancy sequences)
//
// - HaltonSequence<Dim>: Halton sequence (multidimensional extension of van der Corput)
// - SobolSequence<Dim>: Sobol sequence (Gray code + direction-number table)

#ifndef SANGI_QUASI_RANDOM_HPP
#define SANGI_QUASI_RANDOM_HPP

#include <array>
#include <cstdint>
#include <cstddef>
#include <vector>
#include <stdexcept>
#include <bit>

namespace sangi {

// ================================================================
// Internal utilities
// ================================================================

namespace detail {

// First 50 primes (bases for the Halton sequence)
inline constexpr int kPrimes[] = {
     2,   3,   5,   7,  11,  13,  17,  19,  23,  29,
    31,  37,  41,  43,  47,  53,  59,  61,  67,  71,
    73,  79,  83,  89,  97, 101, 103, 107, 109, 113,
   127, 131, 137, 139, 149, 151, 157, 163, 167, 173,
   179, 181, 191, 193, 197, 199, 211, 223, 227, 229
};

// van der Corput sequence: reflect n in base 'base'
inline double vanDerCorput(uint64_t n, int base) {
    double result = 0.0;
    double f = 1.0 / base;
    while (n > 0) {
        result += (n % base) * f;
        n /= base;
        f /= base;
    }
    return result;
}

} // namespace detail

// ================================================================
// HaltonSequence<Dim> — Halton sequence
// ================================================================

// Generates a van der Corput sequence for each dimension with mutually prime bases (2, 3, 5, 7, ...).
// Improves Monte Carlo integration convergence from O(1/N) to O((log N)^D / N).
//
// Limitation: Dim <= 50 (up to 50 dimensions)
template<size_t Dim>
class HaltonSequence {
    static_assert(Dim >= 1 && Dim <= 50, "HaltonSequence: Dim must be 1..50");
public:
    HaltonSequence() : index_(0) {}

    // Start from the specified index
    explicit HaltonSequence(uint64_t startIndex) : index_(startIndex) {}

    // Generate the next point
    std::array<double, Dim> next() {
        std::array<double, Dim> point;
        for (size_t d = 0; d < Dim; ++d)
            point[d] = detail::vanDerCorput(index_, detail::kPrimes[d]);
        ++index_;
        return point;
    }

    // Generate n points at once
    std::vector<std::array<double, Dim>> generate(size_t n) {
        std::vector<std::array<double, Dim>> result;
        result.reserve(n);
        for (size_t i = 0; i < n; ++i)
            result.push_back(next());
        return result;
    }

    // Current index
    uint64_t index() const { return index_; }

    // Reset the index
    void reset(uint64_t startIndex = 0) { index_ = startIndex; }

private:
    uint64_t index_;
};

// ================================================================
// SobolSequence<Dim> — Sobol sequence
// ================================================================

// Gray-code Sobol sequence. Embeds the initial part of the Joe-Kuo direction-number table
// (up to 21201 dimensions).
//
// Limitation: Dim <= 8 (embedded table); for higher dimensions load an external file

namespace detail {

// Direction-number initialization parameters: (degree s, coefficient a, initial direction numbers m[1..s])
// Dimensions 2-8 from the Joe-Kuo table
struct SobolDirParams {
    int s;
    uint32_t a;
    uint32_t m[13]; // m[0..s-1]
};

inline constexpr SobolDirParams kSobolParams[] = {
    // dim 2: s=1, a=0, m={1}
    { 1, 0, {1} },
    // dim 3: s=2, a=1, m={1,1}
    { 2, 1, {1, 1} },
    // dim 4: s=3, a=1, m={1,1,1}
    { 3, 1, {1, 1, 1} },
    // dim 5: s=3, a=2, m={1,3,1}
    { 3, 2, {1, 3, 1} },
    // dim 6: s=4, a=1, m={1,1,1,1}
    { 4, 1, {1, 1, 1, 1} },
    // dim 7: s=4, a=4, m={1,3,5,13}
    { 4, 4, {1, 3, 5, 13} },
    // dim 8: s=5, a=2, m={1,1,5,5,17}
    { 5, 2, {1, 1, 5, 5, 17} },
};

} // namespace detail

template<size_t Dim>
class SobolSequence {
    static_assert(Dim >= 1 && Dim <= 8, "SobolSequence: Dim must be 1..8 (built-in tables)");

    static constexpr int BITS = 52; // number of mantissa bits in a double

public:
    SobolSequence() : index_(0), x_{} {
        initDirectionNumbers();
    }

    // Generate the next point
    std::array<double, Dim> next() {
        if (index_ == 0) {
            ++index_;
            std::array<double, Dim> zero{};
            return zero;
        }

        // Gray code: changing bit position = rightmost zero bit of (index-1)
        int c = std::countr_one(static_cast<uint64_t>(index_ - 1));

        std::array<double, Dim> point;
        double norm = 1.0 / (uint64_t(1) << BITS);
        for (size_t d = 0; d < Dim; ++d) {
            x_[d] ^= v_[d][c];
            point[d] = x_[d] * norm;
        }
        ++index_;
        return point;
    }

    // Generate n points at once
    std::vector<std::array<double, Dim>> generate(size_t n) {
        std::vector<std::array<double, Dim>> result;
        result.reserve(n);
        for (size_t i = 0; i < n; ++i)
            result.push_back(next());
        return result;
    }

    uint64_t index() const { return index_; }

    void reset() {
        index_ = 0;
        x_ = {};
    }

private:
    uint64_t index_;
    std::array<uint64_t, Dim> x_;              // current state
    std::array<std::array<uint64_t, BITS>, Dim> v_; // direction-number table

    void initDirectionNumbers() {
        // Dimension 1 (dim=0): van der Corput (base 2)
        for (int i = 0; i < BITS; ++i)
            v_[0][i] = uint64_t(1) << (BITS - 1 - i);

        // Dimension 2 onward: generated from Joe-Kuo parameters
        for (size_t d = 1; d < Dim; ++d) {
            const auto& par = detail::kSobolParams[d - 1];
            int s = par.s;
            uint32_t a = par.a;

            // Initial direction numbers (normalized to odd)
            for (int i = 0; i < s && i < BITS; ++i)
                v_[d][i] = uint64_t(par.m[i]) << (BITS - 1 - i);

            // Generate the rest via the recurrence
            for (int i = s; i < BITS; ++i) {
                uint64_t vi = v_[d][i - s] ^ (v_[d][i - s] >> s);
                for (int k = 1; k < s; ++k) {
                    if ((a >> (s - 1 - k)) & 1)
                        vi ^= v_[d][i - k];
                }
                v_[d][i] = vi;
            }
        }
    }
};

} // namespace sangi

#endif // SANGI_QUASI_RANDOM_HPP