"""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 sorted(ts) == list(range(3383, 3440, 5))
    assert 4965*3383 % M == 2995
    assert lower == 218767 + (40302 << 18)
    world = (lower + (2995 << 35) - coordinate_term(51069, 419124)) & MASK48
    assert world == 15959335054446
    # Use the film's X-first indexing and check all 25 accepted outputs.
    terms = [(coordinate_term(51069+i%5, 419124+i//5)-coordinate_term(51069, 419124)) & MASK48
             for i in range(25)]
    expected = [1114693760,660285140,2105382940,47575410,1796905160,
                2074383590,1410817530,1734889590,2009625200,440710690,
                1664890290,537178720,1699969350,2115254790,688097620,
                1842255790,1387847170,1302453480,1264531310,344660,
                1757217670,204531470,2132545510,1606583470,569937090]
    for i, term in enumerate(terms):
        output = first_output(lower+(2995 << 35), term)
        assert output == expected[i] and output % 10 == 0 and output < (1 << 31)-8
    # Check every row of the selected candidate group and its two constant pieces.
    root, base = 218767, 40206
    masks = []
    for i, term in enumerate(terms[1:], 1):
        diffs = [(first_output(root+((base+16*j) << 18), term) -
                  first_output(root+((base+16*j) << 18), 0)) % (1 << 31) for j in range(16)]
        masks.append(sum(1 << j for j, value in enumerate(diffs) if value % 10 in (0, 8)))
        if i == 1:
            assert diffs == [1693075028]*8 + [411302484]*8
    assert [masks[i-1] for i in (1,3,4,8,9)] == [0xff,0xfffc,0xfff8,0xfff0,0xffc0]
    combined = 65535
    for mask in masks:
        combined &= mask
    assert combined == 0xc0
    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()