import torch
from safetensors.torch import load_file

def load_model(path='model.safetensors'):
    return load_file(path)

def prefix_and(x3, x2, x1, x0, w):
    """4-bit prefix AND: y_i = AND(x3, x2, ..., x_i)"""
    inp = torch.tensor([float(x3), float(x2), float(x1), float(x0)])

    y3 = int((inp @ w['y3.weight'].T + w['y3.bias'] >= 0).item())
    y2 = int((inp @ w['y2.weight'].T + w['y2.bias'] >= 0).item())
    y1 = int((inp @ w['y1.weight'].T + w['y1.bias'] >= 0).item())
    y0 = int((inp @ w['y0.weight'].T + w['y0.bias'] >= 0).item())

    return y3, y2, y1, y0

if __name__ == '__main__':
    w = load_model()
    print('Prefix-AND Truth Table:')
    print('x3x2x1x0 | y3y2y1y0 | meaning')
    print('---------+----------+--------')
    for i in [0b1111, 0b1110, 0b1101, 0b1011, 0b0111, 0b1100, 0b0000]:
        x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
        y3, y2, y1, y0 = prefix_and(x3, x2, x1, x0, w)
        meaning = "all 1s" if y0 == 1 else f"first 0 at {3 - [y3,y2,y1,y0].index(0) if 0 in [y3,y2,y1,y0] else 'none'}"
        print(f' {x3} {x2} {x1} {x0}  |  {y3} {y2} {y1} {y0}   | {meaning}')