// mars::rng tests — reference vectors + state (de)serialization.
#include <cstdint>
#include <cstdio>
#include <cstring>

#include "mars/rng/mt19937.h"

using mars::rng::MT19937;

static int fails = 0;
#define CHECK(cond)                                                                     \
    do {                                                                                \
        if (!(cond)) {                                                                  \
            std::printf("FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond);                 \
            ++fails;                                                                    \
        }                                                                               \
    } while (0)

int main() {
    // --- standard MT19937 outputs for init_genrand(5489) ------------------------
    {
        MT19937 r(5489u);
        static const uint32_t expect[10] = {3499211612u, 581869302u,  3890346734u, 3586334585u, 545404204u,
                                            4161255391u, 3922919429u, 949333985u,  2715962298u, 1323567403u};
        for (uint32_t e : expect) CHECK(r.next_u32() == e);
        // the 10000th output of the 5489 stream (well-known check value)
        MT19937 r2(5489u);
        uint32_t v = 0;
        for (int i = 0; i < 10000; ++i) v = r2.next_u32();
        CHECK(v == 4123659995u);
    }
    // --- a fresh generator has twisted once: left == N ------------------------
    {
        MT19937 r(1u);
        CHECK(r.left() == MT19937::N);
        CHECK(r.index() == 0);
        r.next_u32();
        CHECK(r.left() == MT19937::N - 1);
        for (int i = 1; i < MT19937::N; ++i) r.next_u32();
        CHECK(r.left() == 0);  // block exhausted; the next draw twists
        r.next_u32();
        CHECK(r.left() == MT19937::N - 1);
    }
    // --- float mapping: y / (2^32 - 1), narrowed to float -----------------------
    // The divisor is 2^32 - 1, so 0xffffffff maps to exactly 1.0 and the range is
    // closed at both ends. The unit value is formed in double and the consumer
    // narrows it to float32.
    {
        CHECK(MT19937::kUnitScale == 1.0 / 4294967295.0);
        CHECK(MT19937::kUnitScale != 1.0 / 4294967296.0);
        CHECK(MT19937::unit_from(0u) == 0.0);
        CHECK(MT19937::unit_from(0xffffffffu) == 1.0);
        CHECK(MT19937::float_from(0xffffffffu) == 1.0f);
        CHECK(MT19937::unit_from(1u) == 1.0 / 4294967295.0);
        // A word with the high bit set exercises the unsigned fix-up path.
        CHECK(MT19937::unit_from(0x80000000u) == 2147483648.0 / 4294967295.0);

        MT19937 r(5489u);
        float f = r.next_float();
        float expect = static_cast<float>(3499211612.0 / 4294967295.0);
        CHECK(f == expect);
        CHECK(f >= 0.f && f <= 1.f);

        // The 2^-32 mapping is wrong but only barely: the two differ by 2^-32 relative,
        // which is far below a float32 ulp, so they disagree for roughly one word in a
        // hundred. Pin that the difference is real but rare, so nobody reads a passing
        // behavioural compare as proof that either divisor would do.
        int differ = 0;
        MT19937 q(17u);
        for (int i = 0; i < 20000; ++i) {
            const uint32_t y = q.next_u32();
            if (static_cast<float>(static_cast<double>(y) * (1.0 / 4294967296.0)) !=
                MT19937::float_from(y))
                ++differ;
        }
        CHECK(differ > 0);
        CHECK(differ < 20000 / 10);

        // next_unit / next_float draw from the same stream position.
        MT19937 a(11u), b(11u);
        CHECK(static_cast<float>(a.next_unit()) == b.next_float());
    }
    // --- the 24-bit-precision variant: same for most words, and never far off ---
    {
        int differ = 0;
        MT19937 r(31u);
        for (int i = 0; i < 20000; ++i) {
            const uint32_t y = r.next_u32();
            const float a = MT19937::float_from(y);
            const float b = MT19937::float_from_pc24(y);
            if (a != b) ++differ;
            const double d = static_cast<double>(a) - static_cast<double>(b);
            CHECK(d < 1e-6 && d > -1e-6);
        }
        // The two agree on the vast majority of words; the point of the check is that
        // the choice of x87 precision mode can only move the last bit.
        CHECK(differ < 20000 / 2);
    }
    // --- next_int_inclusive: mask covers n, the bound is INCLUSIVE --------------
    {
        CHECK(MT19937::cover_mask(0u) == 0u);
        CHECK(MT19937::cover_mask(1u) == 1u);
        CHECK(MT19937::cover_mask(100u) == 127u);
        CHECK(MT19937::cover_mask(255u) == 255u);
        CHECK(MT19937::cover_mask(256u) == 511u);

        MT19937 r(7u);
        bool saw_top = false;
        for (int i = 0; i < 5000; ++i) {
            uint32_t v = r.next_int_inclusive(10);
            CHECK(v <= 10);
            if (v == 10) saw_top = true;
        }
        CHECK(saw_top);  // 10 itself is reachable: the bound is inclusive

        // mask == n means no rejection ever, so exactly one word is consumed
        MT19937 a(9u), b(9u);
        CHECK(a.next_int_inclusive(255u) == (b.next_u32() & 255u));
        CHECK(a.left() == b.left());

        // n == 0 still consumes a word and always yields 0
        MT19937 c(13u), d(13u);
        CHECK(c.next_int_inclusive(0u) == 0u);
        d.next_u32();
        CHECK(c.left() == d.left());
    }
    // --- the compare design: a generator seeded from a pre-call snapshot reads the
    // same stream, and ends in the same state after the same number of draws ---------
    {
        MT19937 live(0xc0ffeeu);
        for (int i = 0; i < 1500; ++i) live.next_u32();  // wherever the game happens to be

        // snapshot "before"
        uint8_t before[MT19937::kStateBytes];
        live.save_state(before);

        // the original consumes some words...
        float a1 = live.next_float();
        float a2 = live.next_float();
        uint8_t after[MT19937::kStateBytes];
        live.save_state(after);

        // ...and ours, seeded from the snapshot, must produce the same values and land on
        // the same state. Matching post-states is the evidence that the draw counts agree.
        MT19937 ours(1u);
        CHECK(ours.load_state(before, sizeof before));
        CHECK(ours.next_float() == a1);
        CHECK(ours.next_float() == a2);
        uint8_t ours_after[MT19937::kStateBytes];
        ours.save_state(ours_after);
        CHECK(std::memcmp(after, ours_after, sizeof after) == 0);

        // one draw too few leaves a different state, so the check has teeth
        MT19937 short_(1u);
        CHECK(short_.load_state(before, sizeof before));
        short_.next_float();
        uint8_t short_after[MT19937::kStateBytes];
        short_.save_state(short_after);
        CHECK(std::memcmp(after, short_after, sizeof after) != 0);
    }
    // --- save_state / load_state round trip, blob layout mt[624] + left --------
    {
        MT19937 r(123456u);
        for (int i = 0; i < 700; ++i) r.next_u32();  // past one twist
        uint8_t blob[MT19937::kStateBytes];
        r.save_state(blob);
        CHECK(MT19937::kStateBytes == 0x9c4);
        uint32_t left_in_blob = uint32_t(blob[2496]) | (uint32_t(blob[2497]) << 8) | (uint32_t(blob[2498]) << 16) |
                                (uint32_t(blob[2499]) << 24);
        CHECK(int(left_in_blob) == r.left());
        CHECK(std::memcmp(blob, r.state(), 4) == 0 || true);  // first word is mt[0] little-endian
        uint32_t w0 = uint32_t(blob[0]) | (uint32_t(blob[1]) << 8) | (uint32_t(blob[2]) << 16) | (uint32_t(blob[3]) << 24);
        CHECK(w0 == r.state()[0]);

        MT19937 s(1u);
        CHECK(s.load_state(blob, sizeof blob));
        CHECK(s.left() == r.left());
        for (int i = 0; i < 2000; ++i) CHECK(s.next_u32() == r.next_u32());

        // truncated / out-of-range blobs are rejected
        CHECK(!s.load_state(blob, sizeof blob - 1));
        uint8_t bad[MT19937::kStateBytes];
        std::memcpy(bad, blob, sizeof bad);
        bad[2496] = 0x71;  // left = 625 > N
        bad[2497] = 0x02;
        CHECK(!s.load_state(bad, sizeof bad));
    }
    // --- load_state(mt, left) positions the next word at mt[N - left] -----------
    {
        MT19937 r(42u);
        uint32_t st[MT19937::N];
        std::memcpy(st, r.state(), sizeof st);
        MT19937 s(0u);
        s.load_state(st, 5);
        for (int i = 0; i < MT19937::N - 5; ++i) r.next_u32();
        for (int i = 0; i < 100; ++i) CHECK(s.next_u32() == r.next_u32());
    }
    // --- range helpers: word cost is the point, not just the value --------------
    // The original exposes these as separate entry points and they cost different
    // numbers of words.  A ledger that counts "draws" without knowing which entry
    // point was called is short by exactly the difference.
    {
        // float_range: EXACTLY ONE word, and it is the same word next_float() would
        // have taken -- so a caller that swaps one for the other keeps the stream.
        MT19937 a(5489u), b(5489u);
        const uint32_t y0 = b.next_u32();
        const float got = a.float_range(-1.0f, 3.0f);
        CHECK(a.left() == MT19937::N - 1);              // one word, not two
        CHECK(got == MT19937::range_from(y0, -1.0f, 3.0f));
        // degenerate span consumes a word all the same
        MT19937 c(7u);
        CHECK(c.float_range(2.5f, 2.5f) == 2.5f);
        CHECK(c.left() == MT19937::N - 1);
        // endpoints: unit is closed at both ends, so both endpoints are reachable
        CHECK(MT19937::range_from(0u, -1.0f, 3.0f) == -1.0f);
        CHECK(MT19937::range_from(0xffffffffu, -1.0f, 3.0f) == 3.0f);
        // the double narrowing is observable: rounding the product to float first is
        // not the same as evaluating the whole expression in double.
        int differs = 0;
        MT19937 d(11u);
        for (int i = 0; i < 20000; ++i) {
            const uint32_t y = d.next_u32();
            const double span = 3.0 - (-1.0);
            const float one_rounding = static_cast<float>(-1.0 + span * MT19937::unit_from(y));
            if (MT19937::range_from(y, -1.0f, 3.0f) != one_rounding) ++differs;
        }
        CHECK(differs > 0);  // if this ever reads 0 the two models are not distinguishable
    }
    {
        // int_range_bell: AT LEAST TWO words, and the split is (h/2, h - h/2) in that
        // order.  h = 7 -> half = 3, so the draws are inclusive [0,3] then [0,4].
        MT19937 a(5489u), b(5489u);
        const int32_t got = a.int_range_bell(10, 17);
        const uint32_t d0 = b.next_int_inclusive(3u);
        const uint32_t d1 = b.next_int_inclusive(4u);
        CHECK(got == int32_t(10u + d0 + d1));
        CHECK(a.index() == b.index());          // same words consumed, same order
        CHECK(a.index() >= 2);                  // never fewer than two
        CHECK(got >= 10 && got <= 17);
        // an empty span still costs two words: half = 0 masks to 0, accepted at once.
        MT19937 c(3u);
        CHECK(c.int_range_bell(4, 4) == 4);
        CHECK(c.index() == 2);
        // and it is triangular, not uniform: the middle of [0,10] must beat the ends.
        int hist[11] = {0};
        MT19937 e(2024u);
        for (int i = 0; i < 40000; ++i) ++hist[e.int_range_bell(0, 10)];
        CHECK(hist[5] > hist[0] * 2);
        CHECK(hist[5] > hist[10] * 2);
    }
    std::printf("test_rng: %s\n", fails ? "FAILED" : "ok");
    return fails ? 1 : 0;
}