"""The 9-step local-collision pattern cancels in the linearised model, and the vectorised step equals the reference."""
import pathlib, random, sys
import numpy as np
ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src")); sys.path.insert(0, str(ROOT / "experiments" / "reduced_round_sha256"))
from wp34 import sha256 as S
import local_collision as L


def test_pattern_cancels_for_all_single_bits_and_random_words():
    for b in range(32):
        assert L.verify_linear(1 << b)[0]
    rng = random.Random(5)
    for _ in range(10):
        assert L.verify_linear(rng.getrandbits(32))[0]


def test_wrong_pattern_does_not_cancel():
    d = 1 << 7
    A, B = L.linear_step_matrices()
    x = np.zeros(256, dtype=np.uint8)
    pat = L.pattern(d); pat[5] = 0  # drop one correction
    for dw in pat:
        x = (A @ x + B @ L.vec(dw)) % 2
    assert x.any()


def test_vectorised_step_matches_reference():
    rng = np.random.default_rng(1)
    st = [rng.integers(0, 2**32, 64, dtype=np.uint32) for _ in range(8)]
    w = rng.integers(0, 2**32, 64, dtype=np.uint32)
    out = L.np_step(st, S.K[3], w)
    for i in range(64):
        ref = S.step(tuple(int(x[i]) for x in st), S.K[3], int(w[i]))
        assert tuple(int(x[i]) for x in out) == ref