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

// IntRandom.hpp — 多倍長乱数整数生成
//
// GMP 互換 API:
//   randomBits(n)       : [0, 2^n) の一様乱数 (mpz_urandomb)
//   randomBelow(upper)  : [0, upper) の一様乱数 (mpz_urandomm)

#ifndef CALX_INT_RANDOM_HPP
#define CALX_INT_RANDOM_HPP

#include <math/core/mp/Int/IntBase.hpp>
#include <random>
#include <cstdint>

namespace calx {

/**
 * @brief 多倍長乱数整数ジェネレータ
 *
 * 内部に mt19937_64 エンジンを保持し、任意ビット幅の一様乱数 Int を生成する。
 * デフォルトコンストラクタは std::random_device でシードする。
 */
class IntRandom {
public:
    /// デフォルトコンストラクタ (random_device でシード)
    IntRandom() {
        std::random_device rd;
        std::seed_seq seq{rd(), rd(), rd(), rd(), rd(), rd(), rd(), rd()};
        m_engine.seed(seq);
    }

    /// 明示的シード
    explicit IntRandom(uint64_t seed) : m_engine(seed) {}

    /// [0, 2^bits) の一様乱数を生成 (GMP mpz_urandomb 相当)
    /// bits == 0 → 0 を返す
    Int randomBits(size_t bits) {
        if (bits == 0) return Int(0);

        size_t nwords = (bits + 63) / 64;
        Int result;
        result.m_words.resize_uninitialized(nwords);

        for (size_t i = 0; i < nwords; ++i) {
            result.m_words[i] = m_engine();
        }

        // 最上位ワードのマスク: bits が 64 の倍数でなければ上位ビットをクリア
        size_t topBits = bits % 64;
        if (topBits != 0) {
            result.m_words[nwords - 1] &= (uint64_t(1) << topBits) - 1;
        }

        // normalize: 上位ゼロワードを除去
        while (!result.m_words.empty() && result.m_words.back() == 0) {
            result.m_words.pop_back();
        }

        result.m_sign = result.m_words.empty() ? 0 : 1;
        result.m_state = NumericState::Normal;
        return result;
    }

    /// [0, upper) の一様乱数を生成 (GMP mpz_urandomm 相当)
    /// rejection sampling: upper.bitLength() ビットの乱数を生成し、
    /// upper 以上なら再試行
    /// upper <= 0 → 0 を返す
    Int randomBelow(const Int& upper) {
        if (upper.getSign() <= 0) return Int(0);

        size_t bits = upper.bitLength();

        // rejection sampling
        for (;;) {
            Int r = randomBits(bits);
            if (r < upper) return r;
        }
    }

    /// [lo, hi] の一様乱数を生成 (両端を含む)
    Int randomRange(const Int& lo, const Int& hi) {
        if (lo > hi) return Int(0);
        if (lo == hi) return lo;
        Int range = hi - lo + Int(1);
        return lo + randomBelow(range);
    }

    /// テスト用の偏った乱数 (GMP mpz_rrandomb 相当)
    /// 長いビットラン (連続する 0/1) を持つ bits ビットの正整数を生成。
    /// MSB は常に 1。乗算・除算のエッジケーステストに有用。
    /// bits == 0 → 0 を返す
    Int rrandomb(size_t bits) {
        if (bits == 0) return Int(0);

        size_t nwords = (bits + 63) / 64;
        Int result;
        result.m_words.resize_uninitialized(nwords);
        for (size_t i = 0; i < nwords; ++i)
            result.m_words[i] = 0;

        // ランダム長のビットランで埋める (GMP 方式)
        // 現在のビット値 (0 or 1) を交互に切り替え
        bool current_bit = true;  // MSB 側から 1 で開始
        size_t pos = bits;        // 残りビット数

        while (pos > 0) {
            // ラン長を幾何分布風に決定 (短いランが多い、時々長い)
            // GMP: gmp_urandomb_ui(state, 0 のとき BITS_PER_MP_LIMB が上限)
            size_t max_run = std::min(pos, static_cast<size_t>(64));
            uint64_t r = m_engine();
            // 下位ビットから連続する 0 の数 + 1 をラン長に (幾何分布の近似)
            size_t run_len;
            if (r == 0) {
                run_len = max_run;
            } else {
                run_len = static_cast<size_t>(std::countr_zero(r)) + 1;
                if (run_len > max_run) run_len = max_run;
            }

            // [pos - run_len, pos) のビットを current_bit で埋める
            if (current_bit) {
                for (size_t i = 0; i < run_len; ++i) {
                    size_t bit_pos = pos - 1 - i;
                    result.m_words[bit_pos / 64] |= uint64_t(1) << (bit_pos % 64);
                }
            }
            // current_bit == false なら既に 0 なのでスキップ

            pos -= run_len;
            current_bit = !current_bit;
        }

        // MSB を確実に 1 にする (bits ビットの数として)
        result.m_words[(bits - 1) / 64] |= uint64_t(1) << ((bits - 1) % 64);

        // 最上位ワードの余剰ビットをクリア
        size_t topBits = bits % 64;
        if (topBits != 0) {
            result.m_words[nwords - 1] &= (uint64_t(1) << topBits) - 1;
        }

        // normalize
        while (!result.m_words.empty() && result.m_words.back() == 0)
            result.m_words.pop_back();

        result.m_sign = result.m_words.empty() ? 0 : 1;
        result.m_state = NumericState::Normal;
        return result;
    }

    /// エンジンへのアクセス (外部カスタマイズ用)
    std::mt19937_64& engine() { return m_engine; }

private:
    std::mt19937_64 m_engine;
};

/// thread_local グローバルエンジンを使うフリー関数

/// [0, 2^bits) の一様乱数
inline Int randomBits(size_t bits) {
    thread_local IntRandom g_rng;
    return g_rng.randomBits(bits);
}

/// [0, upper) の一様乱数
inline Int randomBelow(const Int& upper) {
    thread_local IntRandom g_rng;
    return g_rng.randomBelow(upper);
}

/// [lo, hi] の一様乱数
inline Int randomRange(const Int& lo, const Int& hi) {
    thread_local IntRandom g_rng;
    return g_rng.randomRange(lo, hi);
}

/// テスト用偏り乱数 (長いビットラン)
inline Int rrandomb(size_t bits) {
    thread_local IntRandom g_rng;
    return g_rng.rrandomb(bits);
}

} // namespace calx

#endif // CALX_INT_RANDOM_HPP