import torch
from safetensors.torch import load_file

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

def booth_recode_ref(y):
    """Reference Booth recoding for 4-bit value."""
    bits = [(y >> i) & 1 for i in range(4)]
    prev = [0] + bits[:3]
    result = []
    for i in range(4):
        add_i = (1 - bits[i]) & prev[i]
        sub_i = bits[i] & (1 - prev[i])
        result.append((add_i, sub_i))
    return result

if __name__ == '__main__':
    print('Booth Radix-2 Recoder')
    print('Converts 4-bit multiplier to add/sub signals')
    print()
    print('Value | Recoding')
    for y in range(16):
        codes = booth_recode_ref(y)
        y_signed = y if y < 8 else y - 16
        code_strs = ['+1' if a else ('-1' if s else ' 0') for a, s in codes]
        print(f'  {y_signed:3d} | {" ".join(code_strs)}')