"""Replay the pinned BitNet NumPy packing functions; no torch or CUDA is loaded."""
import ast
import hashlib
import json
from pathlib import Path
import platform
import sys
import numpy as np

path = Path(sys.argv[1])
if path.is_dir():
    path = path / 'gpu' / 'pack_weight.py'
raw = path.read_bytes()
blob = hashlib.sha1(b'blob ' + str(len(raw)).encode() + b'\0' + raw).hexdigest()
assert blob == 'e8194ea0bed82d8a7920e1a598c46540b21d527e', 'source differs from the pin'
names = {'B_global_16x32_to_shared_load_16x32_layout', 'permutate_weight_fastest',
         'compress_int2_to_int8', 'interleave_weight_int8'}
tree = ast.parse(raw.decode())
functions = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
assert {n.name for n in functions} == names
scope = {'np': np}
exec(compile(ast.Module(body=functions, type_ignores=[]), str(path), 'exec'), scope)
layout = scope['B_global_16x32_to_shared_load_16x32_layout']
permute = scope['permutate_weight_fastest']
compress = scope['compress_int2_to_int8']
interleave = scope['interleave_weight_int8']
assert sys.byteorder == 'little', 'the reference decoder in this probe assumes little endian'
print(json.dumps({'python': platform.python_version(), 'numpy': np.__version__, 'byteorder': sys.byteorder}))

mapping = [layout(i, j) for i in range(16) for j in range(32)]
assert len(set(mapping)) == 512
assert set(mapping) == {(i, j) for i in range(16) for j in range(32)}
print(json.dumps({'tileCoordinates': len(mapping), 'uniqueCoordinates': len(set(mapping))}))

def decode_reference(packed):
    # Independent scalar equivalent of decode_i2s_to_i8s: select four byte lanes
    # with 0x03030303, shift two bits per group, then subtract two per lane.
    shape = packed.shape[:-1] + (packed.shape[-1] * 4,)
    output = np.empty(shape, dtype=np.int8)
    words = packed.view(np.uint32)
    for index in np.ndindex(words.shape):
        word = int(words[index])
        base = index[-1] * 16
        for group in range(4):
            selected = (word >> (2 * group)) & 0x03030303
            for lane in range(4):
                output[index[:-1] + (base + group * 4 + lane,)] = ((selected >> (8 * lane)) & 3) - 2
    return output

def restore_tiles(decoded):
    nt, kt = decoded.shape[:2]
    result = np.empty((nt * 16, kt * 32), dtype=np.int8)
    for ni in range(nt):
        for ki in range(kt):
            for i in range(16):
                for j in range(32):
                    row, col = layout(i, j)
                    result[ni * 16 + row, ki * 32 + col] = decoded[ni, ki, i, j]
    return result

rng = np.random.default_rng(7)
cases = 0
for n in [16, 32, 48, 64]:
    for k in [32, 64, 96, 128]:
        for kind in ['zero', 'random']:
            w = np.zeros((n, k), dtype=np.int8) if kind == 'zero' else rng.integers(-1, 2, (n, k), dtype=np.int8)
            tiles = permute(w + 2)
            packed = interleave(compress(tiles), 2)
            restored = restore_tiles(decode_reference(packed))
            np.testing.assert_array_equal(restored, w)
            assert packed.nbytes * 4 == w.nbytes
            cases += 1
print(json.dumps({'roundTripCases': cases, 'reconstructionErrors': 0, 'int8ToPackedByteRatio': 4}))

for n, k in [(16, 33), (17, 32), (16, 31)]:
    w = rng.integers(-1, 2, (n, k), dtype=np.int8)
    tiles = permute(w + 2)
    packed = interleave(compress(tiles), 2)
    # This repeats only the shape arithmetic of the torch wrapper with NumPy.
    # It does not call convert_weight_int8_to_int2 or prove a torch outcome.
    try:
        packed.reshape(n, k // 4)
        reshape = 'accepted'
    except ValueError:
        reshape = 'rejected'
    retained = tiles.size
    print(json.dumps({'shape': [n, k], 'inputValues': w.size, 'retainedValues': retained,
                      'droppedValues': w.size - retained, 'numpyFinalReshape': reshape}))
    if (n, k) == (16, 33):
        assert retained == 512 and reshape == 'accepted'
        np.testing.assert_array_equal(restore_tiles(decode_reference(packed)), w[:, :32])
    elif (n, k) == (17, 32):
        assert retained == 512 and reshape == 'rejected'
    else:
        assert retained == 0 and reshape == 'rejected'
print('PASS: 512-coordinate bijection, 32 round trips, and three shape-boundary fixtures')