"""Check the narrated identities against direct first Java Random outputs.

This checks first-output arithmetic, not exhaustive rejection/redraw coverage.
Run with: uv run python check_lower35.py
"""
import random

A = 0x5DEECE66D
C = A ^ 987234911
MASK48 = (1 << 48) - 1
M = 8192


def first_output(seed, term):
    # Keep the two original initialization XORs separate from the derivation.
    raw_seed = (seed + term) & MASK48
    state = ((raw_seed ^ 987234911) ^ A) & MASK48
    return ((A * state + 11) & MASK48) >> 17


def j32(value):
    return ((value + (1 << 31)) % (1 << 32)) - (1 << 31)


def coordinate_term(x, z):
    return j32(4987142*x*x) + j32(5947611*x) + 4392871*j32(z*z) + j32(389711*z)


def check_case(lower, term):
    q0 = first_output(lower, term)
    r, b = q0 % (1 << 18), q0 >> 18
    for h in range(M):
        t = A*h % M
        direct = first_output(lower + (h << 35), term)
        derived = r + (1 << 18)*((b+t) % M)
        assert direct == derived, (lower, term, h)
        residue = (r-b + (2 if t >= M-b else 0)) % 5
        allowed = (r % 2 == 0 and t % 5 == residue)
        assert allowed == (direct % 10 == 0), (lower, term, h)
        assert 4965*t % M == h
    return r, b


def main():
    assert C.bit_length() == 35
    assert A*4965 % M == 1
    # At a 34-bit boundary C still flips an unknown bit: simple translation fails.
    assert (((1 << 34) ^ C) & MASK48) != ((C+(1 << 34)) & MASK48)
    rng = random.Random(35013)
    cases = [(0, 0), ((1 << 35)-1, -1), (0, MASK48),
             ((1 << 35)-1, (1 << 35)+1)]
    # Real Java-wrapped terms, including signed int limits and negative offsets.
    for x, z in [(51069, 419124), (-1, -1), (2**31-1, -2**31), (0, 0)]:
        cases.append((rng.randrange(1 << 35), coordinate_term(x, z)))
        relative = coordinate_term(x+1 if x < 2**31-1 else x-1, z)-coordinate_term(x, z)
        cases.append((rng.randrange(1 << 35), relative))
    cases += [(rng.randrange(1 << 35), rng.randrange(-2**63, 2**63)) for _ in range(20)]
    for lower, term in cases:
        check_case(lower, term)
    # Find and verify a whole 25-cell intersection, with an actual planted seed.
    lower = 10565146255
    terms = [coordinate_term(51069+dx, 419124+dz)-coordinate_term(51069, 419124)
             for dx in range(5) for dz in range(5)]
    profiles = [(first_output(lower, d) % (1 << 18), first_output(lower, d) >> 18)
                for d in terms]
    points = sorted({0, M, *(M-b for _, b in profiles)})
    ts = []
    if all(r % 2 == 0 for r, _ in profiles):
        for start, stop in zip(points, points[1:]):
            residues = {(r-b+(2 if start >= M-b else 0)) % 5 for r, b in profiles}
            if len(residues) == 1:
                residue = residues.pop()
                ts.extend(range(start+(residue-start) % 5, stop, 5))
    analytic = {4965*t % M for t in ts}
    direct = {h for h in range(M)
              if all(first_output(lower+(h << 35), d) % 10 == 0 for d in terms)}
    assert analytic == direct and direct
    assert 2 + (1 << 18)*((8190+4) % M) == 524290
    assert 4965*4 % M == 3476
    print(f'OK: {len(cases)*M} single-cell identities; '
          f'25-cell interval intersection equals all {M} high-bit trials ({len(direct)} candidates).')


if __name__ == '__main__':
    main()