Finds rng states that produces a given sequence of 128 rand(2) observations when rand(121)s happen after each

sequence_with_jumps.py

"""
Finds rng states that produces a given sequence of 128 rand(2) observations when
rand(121)s happen after each.

This example has a max jump count of 2 (# of times total that rand(121)s reject the result
of rand(128) and cause an additional advancement throughout).

The probability that an arbitrary sequence can be solved for an initial state is roughly
P(X <= max jump count) for X ~ Bin(127, 7/128) (the fact that rand(121) can reject more than once
makes this not exactly correct)

The amount of work needed to solve the sequence scales with CR(127, max jump count)
(combination with replacement; technically sum{0 <= i <= max jump count} CR(127, i) but terms before
the last are comparatively small).

The probability for a max of 2 is ~2.57% and requires checking 8256 combinations.
"""

import numpy as np
import numba
from numba_pokemon_prngs.xorshift import Xoroshiro128PlusRejection

def find_working_example():
    """Find an example sequence & state that has 2 or fewer jumps"""
    rng = Xoroshiro128PlusRejection(0, 0)
    observations = np.empty(128, np.uint8)
    while True:
        jumps = []
        seed_0 = np.random.randint(0, 1 << 64, dtype=np.uint64)
        seed_1 = np.random.randint(0, 1 << 64, dtype=np.uint64)
        rng.re_init(seed_0, seed_1)

        observations[0] = rng.next_rand(2)
        for i in range(1, 128):
            # equivalent to rng.next_rand(121)
            while rng.next_rand(128) >= 121:
                jumps.append(i)
            if len(jumps) > 2:
                break
            observations[i] = rng.next_rand(2)

        if len(jumps) <= 2:
            break

    return seed_0, seed_1, jumps, observations

@numba.njit
def test_sequence(
    rng: Xoroshiro128PlusRejection,
    observations,
    jump_1,
    jump_2
):
    """Test a given combination to see if it produces the expected sequence"""
    # construct a matrix that maps from an initial state to the observations
    mat = np.zeros((128, 128), np.uint8)
    for bit in range(128):
        seed_0 = 1 << bit if bit < 64 else 0
        seed_1 = 1 << (bit - 64) if bit >= 64 else 0
        rng.re_init(seed_0, seed_1)
        mat[bit, 0] = rng.next_rand(2)
        for i in range(1, 128):
            rng.next()
            if jump_1 == i:
                rng.next()
            if jump_2 == i:
                rng.next()
            mat[bit, i] = rng.next_rand(2)

    # invert it to map from observations to the state
    inverse = np.empty((128, 2), np.uint64)
    # identity
    for bit in range(128):
        inverse[bit][0] = 1 << bit if bit < 64 else 0
        inverse[bit][1] = 1 << (bit - 64) if bit >= 64 else 0

    rank = 0
    pivots = []
    for col in range(128):
        for row in range(rank, 128):
            if mat[row, col]:
                for other_row in range(128):
                    if (other_row != row) and mat[other_row, col]:
                        mat[other_row] ^= mat[row]
                        inverse[other_row] ^= inverse[row]
                temp = np.copy(mat[row])
                mat[row] = mat[rank]
                mat[rank] = temp
                temp = np.copy(inverse[row])
                inverse[row] = inverse[rank]
                inverse[rank] = temp
                pivots.append(col)
                rank += 1
                break

    # store the nullbasis in the event that the observations are not determinantive
    nullbasis = np.copy(inverse[rank:])
    # undo pivots
    for i in range(rank - 1, -1, -1):
        pivot = pivots[i]
        temp = np.copy(inverse[i])
        inverse[i] = inverse[pivot]
        inverse[pivot] = temp
    # observations @ inverse
    principal_result = np.zeros(2, np.uint64)
    for i in range(128):
        if observations[i]:
            principal_result ^= inverse[i]

    # loop over other solutions
    for i in range(1 << len(nullbasis)):
        result = np.copy(principal_result)
        for bit in range(128):
            if i == 0:
                break
            if i & 1:
                result ^= nullbasis[bit]
            i >>= 1

        # test result
        rng.re_init(result[0], result[1])
        valid = (rng.next() & 1) == observations[0]
        for i in range(1, 128):
            if not valid:
                break
            rng.next_rand(121)
            valid = (rng.next() & 1) == observations[i]
        if valid:
            print(f"Found state: {result[0]} {result[1]} jumps: {jump_1}, {jump_2}")

@numba.njit
def search(observations):
    """Test all possible combinations of 0-2 jumps"""
    rng = Xoroshiro128PlusRejection(0, 0)
    test_sequence(rng, observations, 0, 0)
    for jump_1 in range(1, 128):
        print(f"{jump_1}/128")
        test_sequence(rng, observations, jump_1, 0)
        for jump_2 in range(jump_1, 128):
            test_sequence(rng, observations, jump_1, jump_2)

if __name__ == "__main__":
    seed_0, seed_1, jumps, observations = find_working_example()
    print(f"Sample state: {seed_0} {seed_1} {jumps=}")
    search(observations)