import torch
from safetensors.torch import load_file

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

def ffs8(bits, w):
    """Find first set bit (1-indexed). Returns 0 if no bits set."""
    inp = torch.tensor([float(b) for b in bits])

    p = []
    for i in range(8):
        pi = int((inp @ w[f'p{i}.weight'].T + w[f'p{i}.bias'] >= 0).item())
        p.append(pi)

    f0 = 1 if (p[0] or p[2] or p[4] or p[6]) else 0
    f1 = 1 if (p[1] or p[2] or p[5] or p[6]) else 0
    f2 = 1 if (p[3] or p[4] or p[5] or p[6]) else 0
    f3 = p[7]

    return f3, f2, f1, f0

if __name__ == '__main__':
    w = load_model()
    print('FFS8 selected tests:')
    print('Input    | FFS | f3f2f1f0')
    print('---------+-----+---------')
    test_vals = [0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x06, 0xF0, 0xFF]
    for val in test_vals:
        bits = [(val >> j) & 1 for j in range(8)]
        f3, f2, f1, f0 = ffs8(bits, w)
        pos = 8*f3 + 4*f2 + 2*f1 + f0
        print(f'{val:08b} |  {pos}  | {f3}{f2}{f1}{f0}')